Audio document RAG

# pip install openai-whisper pyannote.audio torch torchaudio langchain-core langchain-openai cohere python-dotenv

import whisper
from pyannote.audio import Pipeline
import json
from langchain_openai import OpenAIEmbeddings
from langchain_core.vectorstores import InMemoryVectorStore
from langchain.retrievers import MultiQueryRetriever
import cohere
from datetime import datetime
import os
from dotenv import load_dotenv

load_dotenv()

# Step 1: Transcribe audio with timestamps
print("[STEP 1] Transcribing audio...")
whisper_model = whisper.load_model('base')
audio_file = 'interview.mp3'
result = whisper_model.transcribe(audio_file, language='en', verbose=False)

# Step 2: Run speaker diarization
print("[STEP 2] Running speaker diarization...")
diarization_pipeline = Pipeline.from_pretrained(
    'pyannote/speaker-diarization-3.1',
    use_auth_token=os.getenv('HF_TOKEN')
)
diarization = diarization_pipeline(audio_file)

# Step 3: Merge transcription + diarization into speaker turns with metadata
print("[STEP 3] Merging transcript with speaker labels...")
speaker_turns = []
for segment, track, speaker_id in diarization.itertracks(yield_label=True):
    start_ts = segment.start
    end_ts = segment.end
    matching_text = []
    for segment_result in result['segments']:
        if segment_result['start'] >= start_ts and segment_result['end'] <= end_ts:
            matching_text.append(segment_result['text'])
    if matching_text:
        text = ' '.join(matching_text)
        confidence = segment_result.get('confidence', 0.95)
        speaker_turns.append({
            'speaker': f'Speaker_{speaker_id}',
            'text': text.strip(),
            'start_time': round(start_ts, 2),
            'end_time': round(end_ts, 2),
            'confidence': confidence,
            'duration_sec': round(end_ts - start_ts, 2)
        })

print(f"Extracted {len(speaker_turns)} speaker turns")
for turn in speaker_turns[:3]:
    print(f"  {turn['speaker']} [{turn['start_time']}-{turn['end_time']}s, conf={turn['confidence']:.2f}]: {turn['text'][:60]}...")

# Step 4: Create documents for embedding with metadata preservation
print("\n[STEP 4] Creating documents with metadata...")
from langchain_core.documents import Document

documents = []
for idx, turn in enumerate(speaker_turns):
    doc = Document(
        page_content=turn['text'],
        metadata={
            'speaker': turn['speaker'],
            'start_time': turn['start_time'],
            'end_time': turn['end_time'],
            'confidence': turn['confidence'],
            'duration_sec': turn['duration_sec'],
            'source': audio_file,
            'turn_id': idx
        }
    )
    documents.append(doc)

print(f"Created {len(documents)} documents")

# Step 5: Embed documents
print("\n[STEP 5] Embedding documents...")
embeddings = OpenAIEmbeddings(model='text-embedding-3-small')
vectorstore = InMemoryVectorStore.from_documents(
    documents,
    embeddings
)
print(f"Embedded {len(documents)} documents into vector store")

# Step 6: Multi-query retrieval to handle speaker-specific questions
print("\n[STEP 6] Setting up multi-query retriever...")
from langchain_openai import ChatOpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever

llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
multi_query_retriever = MultiQueryRetriever.from_llm_and_query_decomposer(
    vectorstore.as_retriever(search_kwargs={'k': 5}),
    llm_chain=llm,
    prompt_template=None
)

# Step 7: Rerank results by confidence + semantic relevance
print("\n[STEP 7] Configuring reranker...")
co = cohere.ClientV2(api_key=os.getenv('COHERE_API_KEY'))

def rerank_with_confidence(query, docs, top_k=3):
    """Rerank retrieved docs by Cohere score + confidence metadata"""
    doc_texts = [d.page_content for d in docs]
    rerank_result = co.rerank(
        model='rerank-english-v3.0',
        query=query,
        documents=doc_texts,
        top_n=top_k,
        rank_fields=['relevance']
    )
    reranked_docs = []
    for result in rerank_result.results:
        doc = docs[result.index]
        doc.metadata['rerank_score'] = result.relevance_score
        doc.metadata['confidence_weighted_score'] = (
            result.relevance_score * doc.metadata.get('confidence', 0.9)
        )
        reranked_docs.append(doc)
    return sorted(
        reranked_docs,
        key=lambda d: d.metadata['confidence_weighted_score'],
        reverse=True
    )

# Step 8: Example query with speaker context
print("\n[STEP 8] Testing retrieval with speaker query...")
query = "What did Speaker_0 say about the main topic?"
print(f"\nQuery: {query}")

initial_retrieval = multi_query_retriever.invoke(query)
final_results = rerank_with_confidence(query, initial_retrieval, top_k=3)

print(f"\nTop {len(final_results)} results (reranked):")
for i, doc in enumerate(final_results, 1):
    print(f"\n  [{i}] {doc.metadata['speaker']} @ {doc.metadata['start_time']}s (conf={doc.metadata['confidence']:.2f})")
    print(f"      Rerank score: {doc.metadata.get('rerank_score', 0):.3f}")
    print(f"      Text: {doc.page_content[:80]}...")