Update app.py

app.py

@@ -1,94 +1,256 @@
 import gradio as gr
+import librosa
+import soundfile as sf
+import torch
+import json
+import csv
 import os
 import tempfile
-import whisper
+import warnings
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+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
-from transformers import pipeline as hf_pipeline
+import whisper
+
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=RuntimeWarning)
 
 # =========================
-# LOAD MODELS
+# CONFIG
 # =========================
-
 HF_TOKEN = os.getenv("HF_TOKEN")
+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", # Highly recommended to use the latest version
-    token=HF_TOKEN
+    "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"
+)
+sentiment_model = AutoModelForSequenceClassification.from_pretrained(
+    "nlptown/bert-base-multilingual-uncased-sentiment"
+)
+sentiment_model.to(device)
+sentiment_model.eval()
 
-whisper_model = whisper.load_model("base")
+# =========================
+# HELPERS
+# =========================
+def format_time(seconds):
+    s = int(seconds)
+    return f"{s // 60:02d}:{s % 60:02d}"
+
+
+def analyze_sentiment(text):
+    inputs = tokenizer(text[:512], return_tensors="pt", truncation=True).to(device)
+    with torch.no_grad():
+        logits = sentiment_model(**inputs).logits
+    probs = F.softmax(logits, dim=-1)[0]
+    return torch.argmax(probs).item() + 1  # 1–5
 
-sentiment_pipeline = hf_pipeline(
-    "sentiment-analysis",
-    model="nlptown/bert-base-multilingual-uncased-sentiment"
-)
 
 # =========================
-# MAIN FUNCTION
+# MAIN PROCESS
 # =========================
+def process_audio(file, speakers, state):
+    if file is None:
+        return "❌ No audio provided", "", "", state
 
-def analyze_audio(audio_file):
+    temp_wav = None
 
-    if audio_file is None:
-        return "❌ No audio uploaded"
+    try:
+        # Normalise to 16kHz mono WAV (required by Whisper and pyannote)
+        audio, sr = librosa.load(file, sr=16000, mono=True)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
+            sf.write(tmp.name, audio, 16000)
+            temp_wav = tmp.name
 
-    # Save temp file
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
-        temp_path = tmp.name
+        # --- Transcription (Whisper) ---
+        result = whisper_model.transcribe(temp_wav)
+        transcript_text = result["text"].strip()
 
-    os.system(f"ffmpeg -i \"{audio_file}\" -ar 16000 -ac 1 \"{temp_path}\" -y")
+        # --- Diarization (pyannote) ---
+        num_speakers = int(speakers) if speakers > 0 else None
+        diarization = diarization_pipeline(
+            temp_wav,
+            num_speakers=num_speakers
+        )
 
-    # =========================
-    # TRANSCRIPTION
-    # =========================
-    result = whisper_model.transcribe(temp_path)
-    transcript_text = result["text"]
+        # Map raw pyannote speaker IDs → Speaker 1, 2, 3…
+        speaker_map = {}
+        speaker_counter = 1
 
-    # =========================
-    # DIARIZATION
-    # =========================
-    diarization = diarization_pipeline(temp_path)
+        label_map = {
+            1: ("🔴", "Very Negative"),
+            2: ("🟠", "Negative"),
+            3: ("🟡", "Neutral"),
+            4: ("🟢", "Positive"),
+            5: ("🟢", "Very Positive"),
+        }
 
-    # 🔥 MAP RAW SPEAKERS → Speaker 1, 2, 3...
-    speaker_map = {}
-    speaker_counter = 1
+        segments = []
+        conversation = ""
 
-    output = "🎙 TRANSCRIPT + SPEAKERS + SENTIMENT\n\n"
+        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
 
-    for turn, _, speaker in diarization.itertracks(yield_label=True):
+            speaker_id = speaker_map[raw_speaker]
+            start = format_time(turn.start)
+            end = format_time(turn.end)
 
-        # Assign clean speaker labels
-        if speaker not in speaker_map:
-            speaker_map[speaker] = f"Speaker {speaker_counter}"
-            speaker_counter += 1
+            score = analyze_sentiment(transcript_text)
+            emoji, label = label_map.get(score, ("⚪", "Unknown"))
 
-        clean_speaker = speaker_map[speaker]
+            segments.append({
+                "speaker": speaker_id,
+                "start": start,
+                "end": end,
+                "text": transcript_text,
+                "sentiment": label,
+            })
 
-        # Simple text (you can upgrade alignment later)
-        segment_text = transcript_text
+            conversation += (
+                f"Speaker {speaker_id} | Utterance {i}\n"
+                f"({start} - {end})\n"
+                f"{emoji} {label}: {transcript_text}\n\n"
+            )
 
-        sentiment = sentiment_pipeline(segment_text[:512])[0]
+        speaker_count = len(speaker_map)
+        new_state = {"segments": segments, "conversation": conversation}
 
-        output += (
-            f"{clean_speaker} ({turn.start:.2f}s - {turn.end:.2f}s)\n"
-            f"Sentiment: {sentiment['label']} ({round(sentiment['score'],2)})\n"
-            f"Text: {segment_text}\n\n"
+        return (
+            "✅ Done",
+            conversation,
+            f"Speakers: {speaker_count} | Utterances: {len(segments)}",
+            new_state,
         )
 
-    return output
+    except Exception as e:
+        return f"❌ Error: {str(e)}", "", "", state
+
+    finally:
+        if temp_wav and os.path.exists(temp_wav):
+            os.remove(temp_wav)
+
+
+# =========================
+# EXPORT
+# =========================
+def export_file(format_type, state):
+    segments = state.get("segments", [])
+    conversation = state.get("conversation", "")
+
+    if not conversation and not segments:
+        return None
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    if format_type == "TXT":
+        path = f"/tmp/conversation_{timestamp}.txt"
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(conversation)
+
+    elif format_type == "JSON":
+        path = f"/tmp/conversation_{timestamp}.json"
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(segments, f, indent=4)
+
+    elif format_type == "CSV":
+        path = f"/tmp/conversation_{timestamp}.csv"
+        with open(path, "w", newline="", encoding="utf-8") as f:
+            writer = csv.DictWriter(
+                f, fieldnames=["speaker", "start", "end", "text", "sentiment"]
+            )
+            writer.writeheader()
+            writer.writerows(segments)
+
+    elif format_type == "WORD":
+        path = f"/tmp/conversation_{timestamp}.docx"
+        doc = Document()
+        doc.add_heading("Conversation Transcript", 0)
+        doc.add_paragraph(conversation)
+        doc.save(path)
+
+    elif format_type == "PDF":
+        path = f"/tmp/conversation_{timestamp}.pdf"
+        doc = SimpleDocTemplate(path)
+        styles = getSampleStyleSheet()
+        content = [
+            Paragraph(conversation.replace("\n", "<br/>"), styles["Normal"])
+        ]
+        doc.build(content)
+
+    else:
+        return None
+
+    return path
 
 
 # =========================
 # UI
 # =========================
+with gr.Blocks(title="AI Conversation Sentiment Analyzer") as app:
 
-app = gr.Interface(
-    fn=analyze_audio,
-    inputs=gr.Audio(type="filepath", label="Upload Audio"),
-    outputs=gr.Textbox(lines=25, label="Results"),
-    title="🎙 AI Conversation Analyzer",
-    description="Speaker Diarization + Sentiment Analysis"
-)
+    gr.Markdown("# 🎙 AI Conversation Sentiment Analyzer")
+
+    state = gr.State({"segments": [], "conversation": ""})
+
+    with gr.Group():
+        gr.Markdown("### 🎙 Input Audio")
+        audio = gr.Audio(sources=["upload", "microphone"], type="filepath")
+
+    with gr.Group():
+        gr.Markdown("### ⚙ Settings")
+        speakers = gr.Number(value=0, label="Number of speakers (0 = auto-detect)")
+
+    analyze_btn = gr.Button("🚀 Analyze", variant="primary")
+
+    with gr.Group():
+        gr.Markdown("### 💬 Conversation Output")
+        status = gr.Textbox(label="Status")
+        conversation_box = gr.Textbox(lines=18, label="Conversation + Sentiment")
+        info = gr.Textbox(label="Info")
+
+    with gr.Group():
+        gr.Markdown("### 📁 Export")
+        with gr.Row():
+            export_format = gr.Dropdown(
+                ["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],
+        outputs=[status, conversation_box, info, state],
+    )
+
+    export_btn.click(
+        export_file,
+        inputs=[export_format, state],
+        outputs=[download],
+    )
 
-app.launch()
+if __name__ == "__main__":
+    app.launch(theme=gr.themes.Soft())