Update app.py

app.py

@@ -7,6 +7,8 @@ import gradio as gr
 import pandas as pd
 from duckduckgo_search import DDGS
 from transformers import pipeline
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
 
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -18,21 +20,35 @@ class BasicAgent:
         # Initialize the Whisper model for video transcription
         self.whisper_model = whisper.load_model("base")  # You can change the model to `large`, `medium`, etc.
         self.search_pipeline = pipeline("question-answering")
+        self.nlp_model = pipeline("feature-extraction")  # For semantic similarity (using transformer model)
 
-    def call_whisper(self, video_path: str) -> str:
-        # Transcribe the video to text using Whisper model
-        video = moviepy.editor.VideoFileClip(video_path)
-        audio_path = "temp_audio.wav"
-        video.audio.write_audiofile(audio_path)
-        
-        # Transcribe audio to text
-        result = self.whisper_model.transcribe(audio_path)
-        return result["text"]
+    def score_search_results(self, question: str, search_results: list) -> str:
+        # Transform the question and results to embeddings (vector representations)
+        question_embedding = self.nlp_model(question)
+        question_embedding = np.mean(question_embedding[0], axis=0)
+
+        best_score = -1
+        best_answer = None
+
+        # Loop through search results and calculate similarity
+        for result in search_results:
+            result_embedding = self.nlp_model(result['body'])
+            result_embedding = np.mean(result_embedding[0], axis=0)
+
+            # Calculate cosine similarity
+            similarity = cosine_similarity([question_embedding], [result_embedding])
+
+            # Check if this result is better
+            if similarity > best_score:
+                best_score = similarity
+                best_answer = result['body']
+
+        return best_answer
 
     def search(self, question: str) -> str:
-           try:
+        try:
             with DDGS() as ddgs:
-                results = list(ddgs.text(question, max_results=3))
+                results = list(ddgs.text(question, max_results=3))  # Fetch top 3 results
                 if results:
                     # Score all the results and return the best one
                     return self.score_search_results(question, results)
@@ -40,10 +56,16 @@ class BasicAgent:
                     return "No relevant search results found."
         except Exception as e:
             return f"Search error: {e}"
-            
-    def __call__(self, question: str) -> str:
-        # Use the search method to answer the question
-        return self.search(question)
+
+    def call_whisper(self, video_path: str) -> str:
+        # Transcribe the video to text using Whisper model
+        video = moviepy.editor.VideoFileClip(video_path)
+        audio_path = "temp_audio.wav"
+        video.audio.write_audiofile(audio_path)
+        
+        # Transcribe audio to text
+        result = self.whisper_model.transcribe(audio_path)
+        return result["text"]
 
     def __call__(self, question: str, video_path: str = None) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
@@ -194,7 +216,6 @@ with gr.Blocks() as demo:
         1. Please clone this space, then modify the code to define your agent's logic, the tools, the necessary packages, etc ...
         2. Log in to your Hugging Face account using the button below. This uses your HF username for submission.
         3. Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.
-
         ---
         **Disclaimers:**
         Once clicking on the "submit button, it can take quite some time (this is the time for the agent to go through all the questions).