File size: 5,374 Bytes
89d1ea1 de2d5fe 89d1ea1 2eff4c2 207365d 89d1ea1 17ea9f2 89d1ea1 17ea9f2 89d1ea1 17ea9f2 207365d 17ea9f2 89d1ea1 17ea9f2 89d1ea1 17ea9f2 89d1ea1 17ea9f2 89d1ea1 0d9dc0c 89d1ea1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import os
from typing import Optional, Dict, List, Any
import requests
import tempfile
import json
import torch
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}")
final_answer = "Unknown"
try:
if file_name:
search_query = self.extract_search_query(question)
try:
# Set the model ID based on size
model_id = "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 = {}
generate_kwargs["language"] = "english"
# 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:
print(f"Transcription failed: {str(e)}")
raise
except Exception as e:
print(f"Error loading model: {str(e)}")
raise
except Exception as e:
print(f"Transcription failed: {str(e)}")
raise
return final_answer
|