Update app.py

app.py

@@ -36,14 +36,25 @@ aai.settings.api_key = os.getenv("ASSEMBLYAI_API_KEY")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 tokenizer = AutoTokenizer.from_pretrained(
-    "nlptown/bert-base-multilingual-uncased-sentiment"
+    "j-hartmann/emotion-english-distilroberta-base"
 )
 sentiment_model = AutoModelForSequenceClassification.from_pretrained(
-    "nlptown/bert-base-multilingual-uncased-sentiment"
+    "j-hartmann/emotion-english-distilroberta-base"
 )
 sentiment_model.to(device)
 sentiment_model.eval()
 
+# Maps model's 7 emotion classes to business-friendly labels
+EMOTION_LABELS = {
+    0: ("🔴", "Negative"),   # Anger
+    1: ("🔴", "Negative"),   # Disgust
+    2: ("🔴", "Negative"),   # Fear
+    3: ("🟢", "Positive"),   # Joy
+    4: ("🟡", "Neutral"),    # Neutral
+    5: ("🔴", "Negative"),   # Sadness
+    6: ("🟢", "Positive"),   # Surprise
+}
+
 # =========================
 # HELPERS
 # =========================
@@ -52,12 +63,41 @@ def format_time(ms):
     return f"{int(s // 60):02d}:{int(s % 60):02d}"
 
 
+def split_into_chunks(text, chunk_size=200):
+    """
+    Split text into equal fixed-character chunks.
+    Breaks at the nearest space to avoid cutting mid-word.
+    """
+    text = text.strip()
+    if len(text) <= chunk_size:
+        return [text]
+
+    chunks = []
+    while len(text) > chunk_size:
+        split_at = text.rfind(" ", 0, chunk_size)
+        if split_at == -1:
+            split_at = chunk_size
+        chunks.append(text[:split_at].strip())
+        text = text[split_at:].strip()
+
+    if text:
+        chunks.append(text)
+
+    return chunks
+
+
 def analyze_sentiment(text):
-    inputs = tokenizer(text[:512], return_tensors="pt", truncation=True).to(device)
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512,
+        padding=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
+    return torch.argmax(probs).item()
 
 
 def build_segments(transcript):
@@ -96,6 +136,7 @@ def process_audio(file, speakers, language, state):
             speaker_labels=True,
             speakers_expected=int(speakers) if speakers > 0 else None,
             language_code=None if language == "auto" else language,
+            speech_model=aai.SpeechModel.best
         )
         transcript = aai.Transcriber().transcribe(temp_wav, config)
 
@@ -105,26 +146,34 @@ def process_audio(file, speakers, language, state):
         segments = build_segments(transcript)
         speaker_count = len(set(s["speaker"] for s in segments))
 
-        label_map = {
-            1: ("🔴", "Very Negative"),
-            2: ("🟠", "Negative"),
-            3: ("🟡", "Neutral"),
-            4: ("🟢", "Positive"),
-            5: ("🟢", "Very Positive"),
-        }
-
         conversation = ""
-        for i, seg in enumerate(segments, start=1):
-            score = analyze_sentiment(seg["text"])
-            emoji, label = label_map.get(score, ("⚪", "Unknown"))
-            seg["sentiment"] = label
-            conversation += (
-                f"Speaker {seg['speaker']} | Utterance {i}\n"
-                f"({seg['start']} - {seg['end']})\n"
-                f"{emoji} {label}: {seg['text']}\n\n"
-            )
+        export_segments = []
 
-        new_state = {"segments": segments, "conversation": conversation}
+        for i, seg in enumerate(segments, start=1):
+            chunks = split_into_chunks(seg["text"])
+
+            for c_idx, chunk in enumerate(chunks, start=1):
+                emotion_idx = analyze_sentiment(chunk)
+                emoji, label = EMOTION_LABELS.get(emotion_idx, ("⚪", "Unknown"))
+
+                chunk_label = f" | Chunk {c_idx}" if len(chunks) > 1 else ""
+
+                conversation += (
+                    f"Speaker {seg['speaker']} | Utterance {i}{chunk_label}\n"
+                    f"({seg['start']} - {seg['end']})\n"
+                    f"{emoji} {label}: {chunk}\n\n"
+                )
+
+                export_segments.append({
+                    "speaker": seg["speaker"],
+                    "start": seg["start"],
+                    "end": seg["end"],
+                    "chunk": c_idx,
+                    "text": chunk,
+                    "sentiment": label,
+                })
+
+        new_state = {"segments": export_segments, "conversation": conversation}
         return (
             "✅ Done",
             conversation,
@@ -163,7 +212,7 @@ def export_file(format_type, state):
         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"]
+                f, fieldnames=["speaker", "start", "end", "chunk", "text", "sentiment"]
             )
             writer.writeheader()
             writer.writerows(segments)
@@ -235,4 +284,4 @@ with gr.Blocks(title="AI Conversation Sentiment Analyzer", theme=gr.themes.Soft(
     )
 
 if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860)
+    app.launch(server_name="0.0.0.0", server_port=7860)