import streamlit as st import pandas as pd from transcript_extractor import extract_video_id, get_transcript, get_channel_videos, process_videos from data_processor import DataProcessor from database import DatabaseHandler from rag import RAGSystem from query_rewriter import QueryRewriter from evaluation import EvaluationSystem from sentence_transformers import SentenceTransformer import os import json import requests from tqdm import tqdm import sqlite3 # Initialize components @st.cache_resource def init_components(): db_handler = DatabaseHandler() data_processor = DataProcessor() rag_system = RAGSystem(data_processor) query_rewriter = QueryRewriter() evaluation_system = EvaluationSystem(data_processor, db_handler) return db_handler, data_processor, rag_system, query_rewriter, evaluation_system db_handler, data_processor, rag_system, query_rewriter, evaluation_system = init_components() # Ground Truth Generation def generate_questions(transcript): OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'localhost') OLLAMA_PORT = os.getenv('OLLAMA_PORT', '11434') prompt_template = """ You are an AI assistant tasked with generating questions based on a YouTube video transcript. Formulate 10 questions that a user might ask based on the provided transcript. Make the questions specific to the content of the transcript. The questions should be complete and not too short. Use as few words as possible from the transcript. The transcript: {transcript} Provide the output in parsable JSON without using code blocks: {{"questions": ["question1", "question2", ..., "question10"]}} """.strip() prompt = prompt_template.format(transcript=transcript) try: response = requests.post(f'http://{OLLAMA_HOST}:{OLLAMA_PORT}/api/generate', json={ 'model': 'phi3.5', 'prompt': prompt }) response.raise_for_status() return json.loads(response.json()['response']) except requests.RequestException as e: st.error(f"Error generating questions: {str(e)}") return None def generate_ground_truth(video_id): transcript_data = get_transcript(video_id) if transcript_data and 'transcript' in transcript_data: full_transcript = " ".join([entry['text'] for entry in transcript_data['transcript']]) questions = generate_questions(full_transcript) if questions and 'questions' in questions: df = pd.DataFrame([(video_id, q) for q in questions['questions']], columns=['video_id', 'question']) os.makedirs('data', exist_ok=True) df.to_csv('data/ground-truth-retrieval.csv', index=False) st.success("Ground truth data generated and saved to data/ground-truth-retrieval.csv") return df else: st.error("Failed to generate questions.") else: st.error("Failed to generate ground truth data due to transcript retrieval error.") return None # RAG Evaluation def evaluate_rag(sample_size=200): try: ground_truth = pd.read_csv('data/ground-truth-retrieval.csv') except FileNotFoundError: st.error("Ground truth file not found. Please generate ground truth data first.") return None sample = ground_truth.sample(n=min(sample_size, len(ground_truth)), random_state=1) evaluations = [] prompt_template = """ You are an expert evaluator for a Youtube transcript assistant. Your task is to analyze the relevance of the generated answer to the given question. Based on the relevance of the generated answer, you will classify it as "NON_RELEVANT", "PARTLY_RELEVANT", or "RELEVANT". Here is the data for evaluation: Question: {question} Generated Answer: {answer_llm} Please analyze the content and context of the generated answer in relation to the question and provide your evaluation in parsable JSON without using code blocks: {{ "Relevance": "NON_RELEVANT" | "PARTLY_RELEVANT" | "RELEVANT", "Explanation": "[Provide a brief explanation for your evaluation]" }} """.strip() progress_bar = st.progress(0) for i, (_, row) in enumerate(sample.iterrows()): question = row['question'] answer_llm = rag_system.query(question) prompt = prompt_template.format(question=question, answer_llm=answer_llm) evaluation = rag_system.query(prompt) # Assuming rag_system can handle this type of query try: evaluation_json = json.loads(evaluation) evaluations.append((row['video_id'], question, answer_llm, evaluation_json['Relevance'], evaluation_json['Explanation'])) except json.JSONDecodeError: st.warning(f"Failed to parse evaluation for question: {question}") progress_bar.progress((i + 1) / len(sample)) # Store RAG evaluations in the database conn = sqlite3.connect('data/sqlite.db') cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS rag_evaluations ( video_id TEXT, question TEXT, answer TEXT, relevance TEXT, explanation TEXT ) ''') cursor.executemany(''' INSERT INTO rag_evaluations (video_id, question, answer, relevance, explanation) VALUES (?, ?, ?, ?, ?) ''', evaluations) conn.commit() conn.close() st.success("Evaluation complete. Results stored in the database.") return evaluations def main(): st.title("YouTube Transcript RAG System") tab1, tab2, tab3 = st.tabs(["RAG System", "Ground Truth Generation", "Evaluation"]) with tab1: st.header("RAG System") # Input section input_type = st.radio("Select input type:", ["Video URL", "Channel URL", "YouTube ID"]) input_value = st.text_input("Enter the URL or ID:") embedding_model = st.selectbox("Select embedding model:", ["all-MiniLM-L6-v2", "all-mpnet-base-v2"]) if st.button("Process"): with st.spinner("Processing..."): data_processor.embedding_model = SentenceTransformer(embedding_model) if input_type == "Video URL": video_id = extract_video_id(input_value) if video_id: process_single_video(video_id, embedding_model) else: st.error("Failed to extract video ID from the URL") elif input_type == "Channel URL": channel_videos = get_channel_videos(input_value) if channel_videos: process_multiple_videos([video['video_id'] for video in channel_videos], embedding_model) else: st.error("Failed to retrieve videos from the channel") else: process_single_video(input_value, embedding_model) # Query section st.subheader("Query the RAG System") query = st.text_input("Enter your query:") rewrite_method = st.radio("Query rewriting method:", ["None", "Chain of Thought", "ReAct"]) search_method = st.radio("Search method:", ["Hybrid", "Text-only", "Embedding-only"]) if st.button("Search"): with st.spinner("Searching..."): if rewrite_method == "Chain of Thought": query = query_rewriter.rewrite_cot(query) elif rewrite_method == "ReAct": query = query_rewriter.rewrite_react(query) search_method_map = {"Hybrid": "hybrid", "Text-only": "text", "Embedding-only": "embedding"} response = rag_system.query(query, search_method=search_method_map[search_method]) st.write("Response:", response) # Feedback feedback = st.radio("Provide feedback:", ["+1", "-1"]) if st.button("Submit Feedback"): db_handler.add_user_feedback("all_videos", query, 1 if feedback == "+1" else -1) st.success("Feedback submitted successfully!") with tab2: st.header("Ground Truth Generation") video_id = st.text_input("Enter YouTube Video ID for ground truth generation:") if st.button("Generate Ground Truth"): with st.spinner("Generating ground truth..."): ground_truth_df = generate_ground_truth(video_id) if ground_truth_df is not None: st.dataframe(ground_truth_df) csv = ground_truth_df.to_csv(index=False) st.download_button( label="Download Ground Truth CSV", data=csv, file_name="ground_truth.csv", mime="text/csv", ) with tab3: st.header("RAG Evaluation") sample_size = st.number_input("Enter sample size for evaluation:", min_value=1, max_value=1000, value=200) if st.button("Run Evaluation"): with st.spinner("Running evaluation..."): evaluation_results = evaluate_rag(sample_size) if evaluation_results: st.write("Evaluation Results:") st.dataframe(pd.DataFrame(evaluation_results, columns=['Video ID', 'Question', 'Answer', 'Relevance', 'Explanation'])) @st.cache_data def process_single_video(video_id, embedding_model): # Check if the video has already been processed with the current embedding model existing_index = db_handler.get_elasticsearch_index(video_id, embedding_model) if existing_index: st.info(f"Video {video_id} has already been processed with {embedding_model}. Using existing index: {existing_index}") return existing_index transcript_data = get_transcript(video_id) if transcript_data: # Store video metadata in the database video_data = { 'video_id': video_id, 'title': transcript_data['metadata'].get('title', 'Unknown Title'), 'author': transcript_data['metadata'].get('author', 'Unknown Author'), 'upload_date': transcript_data['metadata'].get('upload_date', 'Unknown Date'), 'view_count': int(transcript_data['metadata'].get('view_count', 0)), 'like_count': int(transcript_data['metadata'].get('like_count', 0)), 'comment_count': int(transcript_data['metadata'].get('comment_count', 0)), 'video_duration': transcript_data['metadata'].get('duration', 'Unknown Duration') } db_handler.add_video(video_data) # Store transcript segments in the database for i, segment in enumerate(transcript_data['transcript']): segment_data = { 'segment_id': f"{video_id}_{i}", 'video_id': video_id, 'content': segment.get('text', ''), 'start_time': segment.get('start', 0), 'duration': segment.get('duration', 0) } db_handler.add_transcript_segment(segment_data) # Process transcript for RAG system data_processor.process_transcript(video_id, transcript_data) # Create Elasticsearch index index_name = f"video_{video_id}_{embedding_model}" data_processor.build_index(index_name) # Store Elasticsearch index information db_handler.add_elasticsearch_index(video_id, index_name, embedding_model) st.success(f"Processed and indexed transcript for video {video_id}") st.write("Metadata:", transcript_data['metadata']) return index_name else: st.error(f"Failed to retrieve transcript for video {video_id}") return None @st.cache_data def process_multiple_videos(video_ids, embedding_model): indices = [] for video_id in video_ids: index = process_single_video(video_id, embedding_model) if index: indices.append(index) st.success(f"Processed and indexed transcripts for {len(indices)} videos") return indices if __name__ == "__main__": main()