File size: 1,021 Bytes
b6ed6f7 0aa8e87 b6ed6f7 | 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 | import whisper
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from src.agents.CallState import CallState
class TranscriptionAgent:
def __init__(self):
self.llm = ChatOpenAI(model="gpt-4o", temperature=0)
def __call__(self, state: CallState) -> CallState:
"""Converts audio to text using whisper."""
file_path = state["file_path"]
# Load whisper model - using 'base' for faster processing
model = whisper.load_model("base")
result = model.transcribe(file_path)
state["content"] = result["text"]
# Clean the transcript
prompt = PromptTemplate.from_template(
"Clean the following text of any profanity and fix basic grammatical errors. Return only the clean text:\n{text}"
)
chain = prompt | self.llm
clean_text = chain.invoke({"text": state["content"]}).content
state["clean_content"] = clean_text
return state
|