Update app.py

app.py

@@ -1,280 +1,73 @@
-import os
-import tempfile
-from fastapi import FastAPI, UploadFile, File
+import os, tempfile, uuid
+from fastapi import FastAPI
 import gradio as gr
-import nemo.collections.asr as nemo_asr
-from speechbrain.pretrained import EncoderClassifier
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import soundfile as sf
 import torch
 import numpy as np
-from typing import Dict, List, Tuple
-import json
-import uuid
-from datetime import datetime
+import nemo.collections.asr as nemo_asr
+from speechbrain.pretrained import EncoderClassifier
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Initialize FastAPI app
+# Initialize FastAPI and models
 app = FastAPI()
-
-# Global variables for models
-asr_model = None
-emotion_model = None
-llm_model = None
-llm_tokenizer = None
 conversation_history = {}
 
-def load_models():
-    """Load all required models"""
-    global asr_model, emotion_model, llm_model, llm_tokenizer
-    
-    try:
-        # Load ASR model using correct syntax
-        print("Loading ASR model...")
-        asr_model = nemo_asr.models.ASRModel.from_pretrained("nvidia/parakeet-tdt-0.6b-v2")
-        print("ASR model loaded successfully")
-        
-        # Load emotion recognition model
-        print("Loading emotion model...")
-        emotion_model = EncoderClassifier.from_hparams(
-            source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP",
-            savedir="./emotion_model_cache"
-        )
-        print("Emotion model loaded successfully")
-        
-        # Load LLM for conversation
-        print("Loading LLM...")
-        model_name = "microsoft/DialoGPT-medium"  # Lighter alternative to Vicuna
-        llm_tokenizer = AutoTokenizer.from_pretrained(model_name)
-        llm_model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto" if torch.cuda.is_available() else None
-        )
-        
-        # Add padding token if not present
-        if llm_tokenizer.pad_token is None:
-            llm_tokenizer.pad_token = llm_tokenizer.eos_token
-            
-        print("All models loaded successfully")
-        
-    except Exception as e:
-        print(f"Error loading models: {str(e)}")
-        raise e
-
-def transcribe_audio(audio_path: str) -> Tuple[str, str]:
-    """Transcribe audio and detect emotion"""
-    try:
-        # ASR transcription
-        transcription = asr_model.transcribe([audio_path])
-        text = transcription[0].text if hasattr(transcription[0], 'text') else str(transcription[0])
-        
-        # Emotion detection
-        emotion_result = emotion_model.classify_file(audio_path)
-        emotion = emotion_result[0] if isinstance(emotion_result, list) else str(emotion_result)
-        
-        return text, emotion
-        
-    except Exception as e:
-        print(f"Error in transcription: {str(e)}")
-        return f"Error: {str(e)}", "unknown"
-
-def generate_response(user_text: str, emotion: str, user_id: str) -> str:
-    """Generate contextual response based on user input and emotion"""
-    try:
-        # Get conversation history
-        if user_id not in conversation_history:
-            conversation_history[user_id] = []
-        
-        # Add emotion context to the input
-        emotional_context = f"[User is feeling {emotion}] {user_text}"
-        
-        # Encode input with conversation history
-        conversation_history[user_id].append(emotional_context)
-        
-        # Keep only last 5 exchanges to manage memory
-        if len(conversation_history[user_id]) > 10:
-            conversation_history[user_id] = conversation_history[user_id][-10:]
-        
-        # Create input for the model
-        input_text = " ".join(conversation_history[user_id][-3:])  # Last 3 exchanges
-        
-        # Tokenize and generate
-        inputs = llm_tokenizer.encode(input_text, return_tensors="pt")
-        if torch.cuda.is_available():
-            inputs = inputs.cuda()
-            
-        with torch.no_grad():
-            outputs = llm_model.generate(
-                inputs,
-                max_new_tokens=100,
-                num_return_sequences=1,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=llm_tokenizer.eos_token_id
-            )
-        
-        # Decode response
-        response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract only the new part of the response
-        response = response[len(input_text):].strip()
-        
-        # Add to conversation history
-        conversation_history[user_id].append(response)
-        
-        return response if response else "I understand your feelings. How can I help you today?"
-        
-    except Exception as e:
-        print(f"Error generating response: {str(e)}")
-        return "I'm having trouble processing that right now. Could you try again?"
+# Model loading
+asr_model = nemo_asr.models.ASRModel.from_pretrained("nvidia/parakeet-tdt-0.6b-v2")  # ASR [2]
+emotion_model = EncoderClassifier.from_hparams(
+    source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP",
+    savedir="emotion_cache"
+)  # Emotion [3]
+llm_name = "microsoft/DialoGPT-medium"
+llm_tokenizer = AutoTokenizer.from_pretrained(llm_name)
+llm_model = AutoModelForCausalLM.from_pretrained(llm_name).to("cuda" if torch.cuda.is_available() else "cpu")  # LLM [4]
 
-def process_audio_input(audio_file, user_id: str = None) -> Tuple[str, str, str, str]:
-    """Main processing function for audio input"""
-    if user_id is None:
-        user_id = str(uuid.uuid4())
-    
-    if audio_file is None:
-        return "No audio file provided", "", "", user_id
-    
-    try:
-        # Save uploaded audio to temporary file
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
-            # Handle different audio input formats
-            if hasattr(audio_file, 'name'):
-                # File upload case
-                audio_path = audio_file.name
-            else:
-                # Direct audio data case
-                sf.write(tmp_file.name, audio_file[1], audio_file[0])
-                audio_path = tmp_file.name
-        
-        # Process audio
-        transcription, emotion = transcribe_audio(audio_path)
-        
-        # Generate response
-        response = generate_response(transcription, emotion, user_id)
-        
-        # Clean up temporary file
-        if audio_path != (audio_file.name if hasattr(audio_file, 'name') else ''):
-            os.unlink(audio_path)
-        
-        return transcription, emotion, response, user_id
-        
-    except Exception as e:
-        error_msg = f"Processing error: {str(e)}"
-        print(error_msg)
-        return error_msg, "error", "I'm sorry, I couldn't process your audio.", user_id
+def transcribe_and_emote(audio_path):
+    text = asr_model.transcribe([audio_path])[0].text
+    emotion = emotion_model.classify_file(audio_path)[0]
+    return text, emotion
 
-def get_conversation_history(user_id: str) -> str:
-    """Get formatted conversation history for a user"""
-    if user_id not in conversation_history or not conversation_history[user_id]:
-        return "No conversation history yet."
-    
-    history = conversation_history[user_id]
-    formatted_history = []
-    
-    for i in range(0, len(history), 2):
-        if i + 1 < len(history):
-            user_msg = history[i].replace(f"[User is feeling ", "").split("] ", 1)[-1]
-            bot_msg = history[i + 1]
-            formatted_history.append(f"**You:** {user_msg}")
-            formatted_history.append(f"**AI:** {bot_msg}")
-        
-    return "\n\n".join(formatted_history) if formatted_history else "No conversation history yet."
+def generate_reply(user_text, emotion, uid):
+    # Track and trim history
+    hist = conversation_history.setdefault(uid, [])
+    ctx = f"[Feeling:{emotion}] {user_text}"
+    hist.append(ctx)
+    hist = hist[-6:]
+    conversation_history[uid] = hist
 
-def clear_conversation(user_id: str) -> str:
-    """Clear conversation history for a user"""
-    if user_id in conversation_history:
-        conversation_history[user_id] = []
-    return "Conversation history cleared."
+    prompt = " ".join(hist)
+    inputs = llm_tokenizer.encode(prompt, return_tensors="pt").to(llm_model.device)
+    out = llm_model.generate(inputs, max_new_tokens=100, pad_token_id=llm_tokenizer.eos_token_id)
+    reply = llm_tokenizer.decode(out[0], skip_special_tokens=True)[len(prompt):].strip()
+    hist.append(reply)
+    return reply or "I’m here to help!" 
 
-# Load models on startup
-print("Initializing models...")
-load_models()
-print("Models initialized successfully")
+def process(audio, uid):
+    if not audio:
+        return "", "", "", uid
+    # Save temp file
+    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    data, sr = audio
+    sf.write(tmp.name, data, sr)
+    # ASR + Emotion
+    text, emo = transcribe_and_emote(tmp.name)
+    # LLM response
+    reply = generate_reply(text, emo, uid)
+    # Clean up
+    os.unlink(tmp.name)
+    return text, emo, reply, uid
 
-# Create Gradio interface
-with gr.Blocks(title="Emotional Conversational AI", theme=gr.themes.Soft()) as iface:
-    gr.Markdown("# 🎤 Emotional Conversational AI")
-    gr.Markdown("Upload audio or use your microphone to have an emotional conversation with AI")
-    
-    # User ID state
-    user_id_state = gr.State(value=str(uuid.uuid4()))
-    
-    with gr.Row():
-        with gr.Column(scale=2):
-            # Audio input
-            audio_input = gr.Audio(
-                sources=["microphone", "upload"],
-                type="filepath",
-                label="🎙️ Record or Upload Audio"
-            )
-            
-            # Process button
-            process_btn = gr.Button("🚀 Process Audio", variant="primary", size="lg")
-        
-        with gr.Column(scale=3):
-            # Output displays
-            transcription_output = gr.Textbox(
-                label="📝 Transcription",
-                placeholder="Your speech will appear here...",
-                max_lines=3
-            )
-            
-            emotion_output = gr.Textbox(
-                label="😊 Detected Emotion",
-                placeholder="Detected emotion will appear here...",
-                max_lines=1
-            )
-            
-            response_output = gr.Textbox(
-                label="🤖 AI Response",
-                placeholder="AI response will appear here...",
-                max_lines=5
-            )
-    
-    with gr.Row():
-        with gr.Column():
-            # Conversation history
-            history_output = gr.Textbox(
-                label="💬 Conversation History",
-                placeholder="Your conversation history will appear here...",
-                max_lines=10,
-                interactive=False
-            )
-            
-        with gr.Column():
-            # Control buttons
-            show_history_btn = gr.Button("📖 Show History", variant="secondary")
-            clear_history_btn = gr.Button("🗑️ Clear History", variant="stop")
-            new_session_btn = gr.Button("🆕 New Session", variant="secondary")
-    
-    # Event handlers
-    process_btn.click(
-        fn=process_audio_input,
-        inputs=[audio_input, user_id_state],
-        outputs=[transcription_output, emotion_output, response_output, user_id_state]
-    )
-    
-    show_history_btn.click(
-        fn=get_conversation_history,
-        inputs=[user_id_state],
-        outputs=[history_output]
-    )
-    
-    clear_history_btn.click(
-        fn=clear_conversation,
-        inputs=[user_id_state],
-        outputs=[history_output]
-    )
-    
-    new_session_btn.click(
-        fn=lambda: (str(uuid.uuid4()), "New session started!"),
-        outputs=[user_id_state, history_output]
-    )
+# Gradio interface
+with gr.Blocks() as demo:
+    uid_state = gr.State(value=str(uuid.uuid4()))
+    audio_in = gr.Audio(source="microphone", type="numpy")
+    txt_out = gr.Textbox(label="Transcription")
+    emo_out = gr.Textbox(label="Emotion")
+    rep_out = gr.Textbox(label="AI Reply")
+    btn = gr.Button("Process")
+    btn.click(process, inputs=[audio_in, uid_state], outputs=[txt_out, emo_out, rep_out, uid_state])
 
-# Mount Gradio app to FastAPI
-app = gr.mount_gradio_app(app, iface, path="/")
+app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
     import uvicorn