Update app/generate_ground_truth.py

app/generate_ground_truth.py

@@ -1,15 +1,13 @@
 import pandas as pd
 import json
 from tqdm import tqdm
-import ollama
-from elasticsearch import Elasticsearch
-import sqlite3
 import logging
 import os
-import re
 import sys
+import re
+import sqlite3
 
-# Configure logging for stdout only
+# Configure logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
@@ -18,27 +16,11 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 def extract_model_name(index_name):
-    # Extract the model name from the index name
     match = re.search(r'video_[^_]+_(.+)$', index_name)
     if match:
         return match.group(1)
     return None
 
-def get_transcript_from_elasticsearch(es, index_name, video_id):
-    try:
-        result = es.search(index=index_name, body={
-            "query": {
-                "match": {
-                    "video_id": video_id
-                }
-            }
-        })
-        if result['hits']['hits']:
-            return result['hits']['hits'][0]['_source']['content']
-    except Exception as e:
-        logger.error(f"Error retrieving transcript from Elasticsearch: {str(e)}")
-    return None
-
 def get_transcript_from_sqlite(db_path, video_id):
     try:
         conn = sqlite3.connect(db_path)
@@ -73,13 +55,12 @@ def generate_questions(transcript, max_retries=3):
     retries = 0
 
     while len(all_questions) < 10 and retries < max_retries:
-        prompt = prompt_template.format(transcript=transcript)
         try:
-            response = ollama.chat(
-                model='phi3.5',
-                messages=[{"role": "user", "content": prompt}]
-            )
-            questions = json.loads(response['message']['content'])['questions']
+            model = pipeline("text-generation", model="google/flan-t5-base", device=-1)
+            response = model(prompt_template.format(transcript=transcript), 
+                           max_length=1024, 
+                           num_return_sequences=1)[0]['generated_text']
+            questions = json.loads(response)['questions']
             all_questions.update(questions)
         except Exception as e:
             logger.error(f"Error generating questions: {str(e)}")
@@ -91,19 +72,11 @@ def generate_questions(transcript, max_retries=3):
     return {"questions": list(all_questions)[:10]}
 
 def generate_ground_truth(db_handler, data_processor, video_id):
-    es = Elasticsearch([f'http://{os.getenv("ELASTICSEARCH_HOST", "localhost")}:{os.getenv("ELASTICSEARCH_PORT", "9200")}'])
-    
     # Get existing questions for this video to avoid duplicates
     existing_questions = set(q[1] for q in db_handler.get_ground_truth_by_video(video_id))
     
-    transcript = None
-    index_name = db_handler.get_elasticsearch_index_by_youtube_id(video_id)
-    
-    if index_name:
-        transcript = get_transcript_from_elasticsearch(es, index_name, video_id)
-    
-    if not transcript:
-        transcript = db_handler.get_transcript_content(video_id)
+    # Get transcript from SQLite
+    transcript = get_transcript_from_sqlite(db_handler.db_path, video_id)
     
     if not transcript:
         logger.error(f"Failed to retrieve transcript for video {video_id}")
@@ -141,10 +114,18 @@ def generate_ground_truth(db_handler, data_processor, video_id):
     logger.info(f"Ground truth data saved to {csv_path}")
     return df
 
+def get_evaluation_display_data(video_id=None):
+    """Get evaluation data from CSV file"""
+    try:
+        csv_df = pd.read_csv('data/evaluation_results.csv')
+        if video_id:
+            csv_df = csv_df[csv_df['video_id'] == video_id]
+        return csv_df
+    except FileNotFoundError:
+        return pd.DataFrame()
+
 def get_ground_truth_display_data(db_handler, video_id=None, channel_name=None):
     """Get ground truth data from both database and CSV file"""
-    import pandas as pd
-    
     # Try to get data from database first
     if video_id:
         data = db_handler.get_ground_truth_by_video(video_id)
@@ -203,18 +184,4 @@ def generate_ground_truth_for_all_videos(db_handler, data_processor):
         return df
     else:
         logger.error("Failed to generate questions for any video.")
-        return None
-    
-def get_evaluation_display_data(video_id=None):
-    """Get evaluation data from both database and CSV file"""
-    import pandas as pd
-    
-    # Try to get data from CSV
-    try:
-        csv_df = pd.read_csv('data/evaluation_results.csv')
-        if video_id:
-            csv_df = csv_df[csv_df['video_id'] == video_id]
-    except FileNotFoundError:
-        csv_df = pd.DataFrame()
-    
-    return csv_df
+        return None