Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import assemblyai as aai
 import librosa
 import soundfile as sf
 import torch
@@ -13,8 +14,6 @@ import torch.nn.functional as F
 from docx import Document
 from reportlab.platypus import SimpleDocTemplate, Paragraph
 from reportlab.lib.styles import getSampleStyleSheet
-from pyannote.audio import Pipeline
-import whisper
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", category=UserWarning)
@@ -23,21 +22,9 @@ warnings.filterwarnings("ignore", category=RuntimeWarning)
 # =========================
 # CONFIG
 # =========================
-HF_TOKEN = os.getenv("HF_TOKEN")
+aai.settings.api_key = os.getenv("ASSEMBLYAI_API_KEY")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Whisper for transcription
-whisper_model = whisper.load_model("base")
-
-# Pyannote for speaker diarization
-diarization_pipeline = Pipeline.from_pretrained(
-    "pyannote/speaker-diarization-3.1",
-    use_auth_token=HF_TOKEN
-)
-if device == "cuda":
-    diarization_pipeline.to(torch.device("cuda"))
-
-# BERT for sentiment
 tokenizer = AutoTokenizer.from_pretrained(
     "nlptown/bert-base-multilingual-uncased-sentiment"
 )
@@ -50,9 +37,9 @@ sentiment_model.eval()
 # =========================
 # HELPERS
 # =========================
-def format_time(seconds):
-    s = int(seconds)
-    return f"{s // 60:02d}:{s % 60:02d}"
+def format_time(ms):
+    s = ms / 1000
+    return f"{int(s // 60):02d}:{int(s % 60):02d}"
 
 
 def analyze_sentiment(text):
@@ -63,36 +50,50 @@ def analyze_sentiment(text):
     return torch.argmax(probs).item() + 1  # 1–5
 
 
+def build_segments(transcript):
+    speaker_map = {}
+    counter = 1
+    segments = []
+    for u in transcript.utterances:
+        raw = str(u.speaker)
+        if raw not in speaker_map:
+            speaker_map[raw] = counter
+            counter += 1
+        segments.append({
+            "speaker": speaker_map[raw],
+            "start": format_time(u.start or 0),
+            "end": format_time(u.end or 0),
+            "text": u.text,
+        })
+    return segments
+
+
 # =========================
 # MAIN PROCESS
 # =========================
-def process_audio(file, speakers, state):
+def process_audio(file, speakers, language, state):
     if file is None:
         return "❌ No audio provided", "", "", state
 
     temp_wav = None
-
     try:
-        # Normalise to 16kHz mono WAV (required by Whisper and pyannote)
-        audio, sr = librosa.load(file, sr=16000, mono=True)
+        audio, sr = librosa.load(file, sr=None, mono=True)
         with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
-            sf.write(tmp.name, audio, 16000)
+            sf.write(tmp.name, audio, sr)
             temp_wav = tmp.name
 
-        # --- Transcription (Whisper) ---
-        result = whisper_model.transcribe(temp_wav)
-        transcript_text = result["text"].strip()
-
-        # --- Diarization (pyannote) ---
-        num_speakers = int(speakers) if speakers > 0 else None
-        diarization = diarization_pipeline(
-            temp_wav,
-            num_speakers=num_speakers
+        config = aai.TranscriptionConfig(
+            speaker_labels=True,
+            speakers_expected=int(speakers) if speakers > 0 else None,
+            language_code=None if language == "auto" else language,
         )
+        transcript = aai.Transcriber().transcribe(temp_wav, config)
 
-        # Map raw pyannote speaker IDs → Speaker 1, 2, 3…
-        speaker_map = {}
-        speaker_counter = 1
+        if transcript.error:
+            return f"❌ {transcript.error}", "", "", state
+
+        segments = build_segments(transcript)
+        speaker_count = len(set(s["speaker"] for s in segments))
 
         label_map = {
             1: ("🔴", "Very Negative"),
@@ -102,40 +103,18 @@ def process_audio(file, speakers, state):
             5: ("🟢", "Very Positive"),
         }
 
-        segments = []
         conversation = ""
-
-        for i, (turn, _, raw_speaker) in enumerate(
-            diarization.itertracks(yield_label=True), start=1
-        ):
-            if raw_speaker not in speaker_map:
-                speaker_map[raw_speaker] = speaker_counter
-                speaker_counter += 1
-
-            speaker_id = speaker_map[raw_speaker]
-            start = format_time(turn.start)
-            end = format_time(turn.end)
-
-            score = analyze_sentiment(transcript_text)
+        for i, seg in enumerate(segments, start=1):
+            score = analyze_sentiment(seg["text"])
             emoji, label = label_map.get(score, ("⚪", "Unknown"))
-
-            segments.append({
-                "speaker": speaker_id,
-                "start": start,
-                "end": end,
-                "text": transcript_text,
-                "sentiment": label,
-            })
-
+            seg["sentiment"] = label
             conversation += (
-                f"Speaker {speaker_id} | Utterance {i}\n"
-                f"({start} - {end})\n"
-                f"{emoji} {label}: {transcript_text}\n\n"
+                f"Speaker {seg['speaker']} | Utterance {i}\n"
+                f"({seg['start']} - {seg['end']})\n"
+                f"{emoji} {label}: {seg['text']}\n\n"
             )
 
-        speaker_count = len(speaker_map)
         new_state = {"segments": segments, "conversation": conversation}
-
         return (
             "✅ Done",
             conversation,
@@ -157,7 +136,6 @@ def process_audio(file, speakers, state):
 def export_file(format_type, state):
     segments = state.get("segments", [])
     conversation = state.get("conversation", "")
-
     if not conversation and not segments:
         return None
 
@@ -193,9 +171,7 @@ def export_file(format_type, state):
         path = f"/tmp/conversation_{timestamp}.pdf"
         doc = SimpleDocTemplate(path)
         styles = getSampleStyleSheet()
-        content = [
-            Paragraph(conversation.replace("\n", "<br/>"), styles["Normal"])
-        ]
+        content = [Paragraph(conversation.replace("\n", "<br/>"), styles["Normal"])]
         doc.build(content)
 
     else:
@@ -219,7 +195,11 @@ with gr.Blocks(title="AI Conversation Sentiment Analyzer") as app:
 
     with gr.Group():
         gr.Markdown("### ⚙ Settings")
-        speakers = gr.Number(value=0, label="Number of speakers (0 = auto-detect)")
+        with gr.Row():
+            speakers = gr.Number(value=0, label="Speakers (0 = auto-detect)")
+            language = gr.Dropdown(
+                ["auto", "en", "fr", "es", "de"], value="auto", label="Language"
+            )
 
     analyze_btn = gr.Button("🚀 Analyze", variant="primary")
 
@@ -233,19 +213,16 @@ with gr.Blocks(title="AI Conversation Sentiment Analyzer") as app:
         gr.Markdown("### 📁 Export")
         with gr.Row():
             export_format = gr.Dropdown(
-                ["TXT", "JSON", "CSV", "WORD", "PDF"],
-                value="TXT",
-                label="Format"
+                ["TXT", "JSON", "CSV", "WORD", "PDF"], value="TXT", label="Format"
             )
             export_btn = gr.Button("⬇ Export")
         download = gr.File()
 
     analyze_btn.click(
         process_audio,
-        inputs=[audio, speakers, state],
+        inputs=[audio, speakers, language, state],
         outputs=[status, conversation_box, info, state],
     )
-
     export_btn.click(
         export_file,
         inputs=[export_format, state],