Update app.py

app.py

@@ -1,272 +1,93 @@
 import gradio as gr
-from dotenv import load_dotenv
-import assemblyai as aai
-import librosa
-import soundfile as sf
-import torch
-import json
-import csv
 import os
 import tempfile
-import time
-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
+import whisper
+from pyannote.audio import Pipeline
+from transformers import pipeline as hf_pipeline
 
 # =========================
-# CONFIG
+# LOAD MODELS
 # =========================
-load_dotenv()
 
-aai.settings.api_key = os.getenv("ASSEMBLYAI_API_KEY")
-hf_token = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-if not hf_token:
-    raise ValueError("❌ HF_TOKEN is missing. Add it to your .env")
+diarization_pipeline = Pipeline.from_pretrained(
+    "pyannote/speaker-diarization",
+    use_auth_token=HF_TOKEN
+)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+whisper_model = whisper.load_model("base")
 
-MODEL_NAME = "nlptown/bert-base-multilingual-uncased-sentiment"
+sentiment_pipeline = hf_pipeline(
+    "sentiment-analysis",
+    model="nlptown/bert-base-multilingual-uncased-sentiment"
+)
 
 # =========================
-# LOAD MODEL (FIXED)
+# MAIN FUNCTION
 # =========================
-def load_hf_model():
-    for attempt in range(3):
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(
-                MODEL_NAME,
-                token=hf_token,
-                cache_dir="./models"
-            )
 
-            model = AutoModelForSequenceClassification.from_pretrained(
-                MODEL_NAME,
-                token=hf_token,
-                cache_dir="./models"
-            )
+def analyze_audio(audio_file):
 
-            return tokenizer, model
+    if audio_file is None:
+        return "❌ No audio uploaded"
 
-        except Exception as e:
-            print(f"⚠️ HF load failed (attempt {attempt+1}): {e}")
-            time.sleep(5)
+    # Save temp file
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
+        temp_path = tmp.name
 
-    raise RuntimeError("❌ Failed to load Hugging Face model after retries")
+    os.system(f"ffmpeg -i \"{audio_file}\" -ar 16000 -ac 1 \"{temp_path}\" -y")
 
-tokenizer, model = load_hf_model()
+    # =========================
+    # TRANSCRIPTION
+    # =========================
+    result = whisper_model.transcribe(temp_path)
+    transcript_text = result["text"]
 
-model.to(device)
-model.eval()
+    # =========================
+    # DIARIZATION
+    # =========================
+    diarization = diarization_pipeline(temp_path)
 
-# =========================
-# GLOBAL
-# =========================
-global_segments = []
-global_conversation = ""
-
-# =========================
-# HELPERS
-# =========================
-def format_time(ms):
-    s = ms / 1000
-    return f"{int(s//60):02d}:{int(s%60):02d}"
-
-
-def build_segments(transcript):
+    # 🔥 MAP RAW SPEAKERS → Speaker 1, 2, 3...
     speaker_map = {}
-    current_id = 1
-    segments = []
-
-    for u in transcript.utterances:
-        raw = str(u.speaker)
-
-        if raw not in speaker_map:
-            speaker_map[raw] = current_id
-            current_id += 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
+    speaker_counter = 1
 
+    output = "🎙 TRANSCRIPT + SPEAKERS + SENTIMENT\n\n"
 
-def analyze_text(text):
-    inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        truncation=True,
-        max_length=512
-    ).to(device)
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
 
-    with torch.no_grad():
-        logits = model(**inputs).logits
+        # Assign clean speaker labels
+        if speaker not in speaker_map:
+            speaker_map[speaker] = f"Speaker {speaker_counter}"
+            speaker_counter += 1
 
-    probs = F.softmax(logits, dim=-1)[0]
-    return torch.argmax(probs).item() + 1
+        clean_speaker = speaker_map[speaker]
 
+        # Simple text (you can upgrade alignment later)
+        segment_text = transcript_text
 
-# =========================
-# MAIN PROCESS
-# =========================
-def process_audio(file, speakers=0, language="auto"):
-    global global_segments, global_conversation
-
-    if file is None:
-        return "❌ No audio provided", "", ""
-
-    path = file if isinstance(file, str) else file.name
-    temp_path = None
-
-    try:
-        # Load audio
-        audio, sr = librosa.load(path, sr=None, mono=True)
-
-        # Create TEMP FILE
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
-            sf.write(tmp.name, audio, sr)
-            temp_path = tmp.name
+        sentiment = sentiment_pipeline(segment_text[:512])[0]
 
-        config = aai.TranscriptionConfig(
-            speaker_labels=True,
-            speakers_expected=speakers if speakers > 0 else None,
-            language_code=None if language == "auto" else language
+        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"
         )
 
-        transcript = aai.Transcriber().transcribe(temp_path, config)
-
-        if transcript.error:
-            return f"❌ {transcript.error}", "", ""
-
-        global_segments = build_segments(transcript)
-        speaker_count = len(set(s["speaker"] for s in global_segments))
-
-        label_map = {
-            1: ("🔴", "Very Negative"),
-            2: ("🟠", "Negative"),
-            3: ("🟡", "Neutral"),
-            4: ("🟢", "Positive"),
-            5: ("🟢", "Very Positive")
-        }
-
-        conversation = ""
-
-        for i, seg in enumerate(global_segments, start=1):
-            score = analyze_text(seg["text"])
-            emoji, label = label_map.get(score, ("⚪", "Unknown"))
-
-            conversation += (
-                f"Speaker {seg['speaker']} | Utterance {i}\n"
-                f"({seg['start']} - {seg['end']})\n"
-                f"{emoji} {label}: {seg['text']}\n\n"
-            )
-
-        global_conversation = conversation
-
-        return "✅ Done", conversation, f"Speakers: {speaker_count} | Utterances: {len(global_segments)}"
-
-    except Exception as e:
-        return f"❌ Error: {str(e)}", "", ""
-
-    finally:
-        if temp_path and os.path.exists(temp_path):
-            os.remove(temp_path)
-
-
-# =========================
-# EXPORT
-# =========================
-def export_file(format_type):
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-
-    if format_type == "TXT":
-        path = f"conversation_{timestamp}.txt"
-        with open(path, "w", encoding="utf-8") as f:
-            f.write(global_conversation)
-
-    elif format_type == "JSON":
-        path = f"conversation_{timestamp}.json"
-        with open(path, "w", encoding="utf-8") as f:
-            json.dump(global_segments, f, indent=4)
-
-    elif format_type == "CSV":
-        path = f"conversation_{timestamp}.csv"
-        with open(path, "w", newline="", encoding="utf-8") as f:
-            writer = csv.DictWriter(f, fieldnames=["speaker", "start", "end", "text"])
-            writer.writeheader()
-            writer.writerows(global_segments)
-
-    elif format_type == "WORD":
-        path = f"conversation_{timestamp}.docx"
-        doc = Document()
-        doc.add_heading("Conversation Transcript", 0)
-        doc.add_paragraph(global_conversation)
-        doc.save(path)
-
-    elif format_type == "PDF":
-        path = f"conversation_{timestamp}.pdf"
-        doc = SimpleDocTemplate(path)
-        styles = getSampleStyleSheet()
-        content = [Paragraph(global_conversation.replace("\n", "<br/>"), styles["Normal"])]
-        doc.build(content)
-
-    return path
+    return output
 
 
 # =========================
 # UI
 # =========================
-with gr.Blocks(title="AI Conversation Sentiment System") as app:
-
-    gr.Markdown("# 🎙 AI Conversation Sentiment Analyzer")
-
-    with gr.Group():
-        gr.Markdown("### 🎙 Input Audio")
-        audio = gr.Audio(sources=["upload", "microphone"], type="filepath")
-
-    with gr.Group():
-        gr.Markdown("### ⚙ Settings")
-        with gr.Row():
-            speakers = gr.Number(value=0, label="Speakers (0 = auto)")
-            language = gr.Dropdown(["auto", "en", "fr", "es", "de"], value="auto")
-
-    analyze_btn = gr.Button("🚀 Analyze")
-
-    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="Select Format"
-            )
-            export_btn = gr.Button("⬇ Export")
-        download = gr.File()
-
-    analyze_btn.click(
-        process_audio,
-        inputs=[audio, speakers, language],
-        outputs=[status, conversation_box, info]
-    )
 
-    export_btn.click(
-        export_file,
-        inputs=[export_format],
-        outputs=[download]
-    )
+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"
+)
 
-if __name__ == "__main__":
-    app.launch(theme=gr.themes.Soft())
+app.launch()