Create audio_agent.py

audio_agent.py

@@ -0,0 +1,111 @@
+import os
+from typing import Optional, Dict, List, Any
+import requests
+import tempfile
+import json
+import whisper
+from base_agent import BaseAgent
+from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
+
+class AudioAgent(BaseAgent):
+    """
+    Agent for transcribing audio/mp3 file.
+    """
+    def __init__(self, model: str):
+        super().__init__(model=model)
+
+    def __call__(self, question: str, file_name: str, **kwargs) -> str:
+        print(f"AudioAgent: Transcribing file: {file_name")
+
+        search_query = self.extract_search_query(question)
+
+        try:
+            # Set the model ID based on size
+            model_id = f"openai/whisper-{tiny}"
+            
+            # Load model with optimizations for CPU
+            model = AutoModelForSpeechSeq2Seq.from_pretrained(
+                model_id,
+                torch_dtype=torch.float32,  # Use float32 for CPU
+                low_cpu_mem_usage=True,
+                use_safetensors=True
+            )
+            
+            # Load processor
+            processor = AutoProcessor.from_pretrained(model_id)
+            
+            # Create pipeline with small batch size for CPU
+            pipe = pipeline(
+                "automatic-speech-recognition",
+                model=model,
+                tokenizer=processor.tokenizer,
+                feature_extractor=processor.feature_extractor,
+                max_new_tokens=128,
+                chunk_length_s=30,
+                batch_size=1,  # Small batch size for CPU
+                return_timestamps=True,
+                device="cpu",
+            )
+            
+            print("Model loaded successfully.")
+
+            print(f"Starting transcription of {file_name}")
+
+            try:
+                print("Running transcription (this may take some time on CPU)...")
+                
+                # Set task and language options
+                generate_kwargs = {}
+                if language:
+                    generate_kwargs["language"] = "auto"
+                    
+                # Run the pipeline directly on the file
+                result = pipe(
+                    file_name,
+                    generate_kwargs=generate_kwargs,
+                    return_timestamps=True
+                )
+                
+                # Process results based on output format
+
+                transcript = result["text"]
+
+                system_prompt = """
+                You are a research expert tasked with answering factual questions based on a text transcript extracted from an audio file.
+                Analyze the transcript carefully and extract the information needed to answer the question accurately.
+                Provide a concise, accurate answer based only on the provided information.
+                You must provide only the exact answer requested by the user, in the correct type and minimal format, based purely on the provided information.
+                YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
+                If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
+                If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
+                If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+                If you cannot find the answer based on the provided information, reply only with "Unknown".
+                """
+                
+                user_prompt = f"""
+                Question: {question}
+                
+                Transcript:
+                {transcript}
+                
+                Based on these search results, provide a concise and accurate answer to the question.
+                """
+
+                final_answer = self.llm_call(system_prompt, user_prompt)
+
+                print(f"Final answer: {final_answer}")
+        
+                if final_answer.lower() == "unknown":
+                    final_answer = self.ask_james(question)                
+
+            except Exception as e:
+                logger.error(f"Transcription failed: {str(e)}")
+                raise
+            
+        except Exception as e:
+            logger.error(f"Error loading model: {str(e)}")
+            raise
+
+    return final_answer
+        
+