helo-ayush/Diarization_VoiceFingerprinted

1# ==============================================================================
2# SARVAM AI PROCESSOR
3# Handles Speech-to-Text and Diarization (speaker separation) via Sarvam AI.
4# Optimized for Indian languages (Hindi, Hinglish, English, etc.)
5# ==============================================================================
6import os
7import time
8import tempfile
9from sarvamai import SarvamAI
10
11client = SarvamAI(api_subscription_key=os.getenv("SARVAM_API_KEY"))
12
13
14async def transcribe_with_sarvam(audio_bytes: bytes, filename: str = "audio.ogg") -> dict:
15    """
16    Transcribe audio with speaker diarization using Sarvam AI's Batch API.
17    
18    Uses the saaras:v3 model optimized for Indian languages (Hindi, Hinglish, etc.)
19    with speaker diarization enabled.
20    """
21    size_kb = len(audio_bytes) / 1024
22    print(f"🎙️ Sending {size_kb:.1f}KB to Sarvam AI...")
23
24    start_time = time.time()
25
26    # Save audio to temp file (Sarvam SDK needs file paths)
27    # The input audio_bytes is already converted to OGG by audio_processor.py
28    with tempfile.NamedTemporaryFile(suffix=".ogg", delete=False) as tmp:
29        tmp.write(audio_bytes)
30        tmp_path = tmp.name
31
32    try:
33        # Create batch job with diarization
34        print("   📋 Creating Sarvam batch job...")
35        job = client.speech_to_text_job.create_job(
36            model="saaras:v3",
37            mode="translit",
38            language_code="hi-IN",
39            with_diarization=True,
40            num_speakers=2,
41        )
42        print(f"   ✅ Job created: {job.id if hasattr(job, 'id') else 'OK'}")
43
44        # Upload audio file
45        print("   📤 Uploading audio file...")
46        job.upload_files(file_paths=[tmp_path])
47
48        # Start processing
49        print("   ⏳ Processing started...")
50        job.start()
51
52        # Wait for completion
53        print("   ⏳ Waiting for completion...")
54        job.wait_until_complete()
55
56        processing_time = int((time.time() - start_time) * 1000)
57        print(f"   ✅ Sarvam completed in {processing_time}ms")
58
59        # Get results
60        file_results = job.get_file_results()
61        successful = file_results.get("successful", [])
62
63        if not successful:
64            failed = file_results.get("failed", [])
65            error_msg = failed[0].get("error_message", "Unknown error") if failed else "No results"
66            print(f"   ❌ Sarvam transcription failed: {error_msg}")
67            return {
68                "words": [],
69                "transcript": "",
70                "has_speech": False,
71                "processing_time": processing_time,
72            }
73
74        # Download and parse the output
75        output_dir = tempfile.mkdtemp()
76        job.download_outputs(output_dir=output_dir)
77
78        # Read the output JSON file
79        import json
80        output_files = os.listdir(output_dir)
81        result_data = None
82        for f in output_files:
83            if f.endswith(".json"):
84                with open(os.path.join(output_dir, f), encoding="utf-8") as fp:
85                    result_data = json.load(fp)
86                break
87
88        if not result_data:
89            print("   ❌ No output file found")
90            return {
91                "words": [],
92                "transcript": "",
93                "has_speech": False,
94                "processing_time": processing_time,
95            }
96
97        # Extract diarized transcript
98        diarized = result_data.get("diarized_transcript", {})
99        entries = diarized.get("entries", [])
100
101        if not entries:
102            # Fallback to plain transcript
103            plain = result_data.get("transcript", "")
104            return {
105                "words": [],
106                "transcript": plain,
107                "has_speech": bool(plain),
108                "processing_time": processing_time,
109            }
110
111        # Build word-level-like data from diarized entries
112        # Sarvam returns sentence-level segments, not word-level
113        words = []
114        for entry in entries:
115            text = entry.get("transcript", "")
116            speaker = int(entry.get("speaker_id", 0))
117            start = entry.get("start_time_seconds", 0)
118            end = entry.get("end_time_seconds", 0)
119
120            # Split into words for compatibility with our pipeline
121            for word in text.split():
122                words.append({
123                    "word": word,
124                    "speaker": speaker,
125                    "confidence": 0.8,  # Sarvam doesn't provide word-level confidence
126                    "start": start,
127                    "end": end,
128                })
129
130        print(f"   📊 Sarvam: {len(entries)} segments, {len(words)} words")
131        return {"words": words, "has_speech": True, "processing_time": processing_time}
132
133    finally:
134        # Cleanup
135        try:
136            os.unlink(tmp_path)
137        except OSError:
138            pass
139