Update app.py

app.py

@@ -1,66 +1,54 @@
+from transformers import pipeline, Wav2Vec2ForCTC, Wav2Vec2Processor
 import gradio as gr
 import torch
 import numpy as np
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Set device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load Wav2Vec2 model and processor for speech recognition
+# Load models
 processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
-model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h").to(device)
+model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h").to("cpu")
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
-# Load FLAN-T5 model for personality generation
-gen_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-gen_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base").to(device)
+labels = [
+    "secret agent", "lecturer", "motivational speaker", "comedian",
+    "philosopher", "calm leader", "emotional artist", "enthusiastic teacher",
+    "strict officer", "mysterious thinker"
+]
 
-# Function to transcribe audio to text
-def transcribe(audio):
+# Transcription + personality
+def transcribe_and_analyze(audio):
     if audio is None:
-        return "Please upload or record an audio file."
-    
-    if isinstance(audio, tuple):
-        audio_np = audio[1]
-    else:
-        audio_np = audio
+        return "No audio provided", "No analysis"
 
-    if isinstance(audio_np, np.ndarray) and audio_np.ndim > 1:
-        audio_np = np.mean(audio_np, axis=1)
-
-    input_values = processor(audio_np, sampling_rate=16000, return_tensors="pt").input_values.to(device)
+    # Mono audio
+    if len(audio.shape) > 1:
+        audio = np.mean(audio, axis=1)
+    
+    inputs = processor(audio, sampling_rate=16000, return_tensors="pt")
+    input_values = inputs.input_values
     with torch.no_grad():
         logits = model(input_values).logits
     predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.decode(predicted_ids[0])
-    return transcription.lower()
+    transcription = processor.decode(predicted_ids[0]).lower()
 
-# Function to generate personality from transcription
-def generate_persona_from_text(transcription):
-    prompt = f"Describe the speaker's personality and role as if they are a fictional character, based on this message:\n\"{transcription}\""
-    inputs = gen_tokenizer(prompt, return_tensors="pt").to(device)
-    output_ids = gen_model.generate(**inputs, max_length=100)
-    return gen_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    result = classifier(transcription, candidate_labels=labels)
+    top_label = result["labels"][0]
+    score = result["scores"][0]
 
-# Complete function for Gradio
-def analyze_speaker(audio):
-    transcription = transcribe(audio)
-    if "please upload" in transcription:
-        return transcription, ""
-    persona = generate_persona_from_text(transcription)
-    return transcription, persona
+    return transcription, f"You sound like a **{top_label}** (confidence: {score:.2f})"
 
-# Gradio Interface
+# Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as app:
-    gr.Markdown("# Voice2Persona AI")
-    gr.Markdown("Upload or record your voice. We'll transcribe it and guess your fictional personality.")
+    gr.Markdown("# Voice2PersonaAI (Free AI-Based Personality Engine)")
+
+    with gr.Row():
+        audio_input = gr.Audio(sources=["microphone", "upload"], type="numpy", label="🎤 Your Voice")
+
+    transcribe_button = gr.Button("Transcribe & Analyze")
 
     with gr.Row():
-    audio_input = gr.Audio(sources=["microphone", "upload"], type="numpy", label="🎤 Your Voice")
-    transcribed_text = gr.Textbox(label="📝 Transcription")
-    persona_output = gr.Textbox(label="🧠 Persona Analysis")
-    analyze_button = gr.Button("Analyze")
+        transcribed_text = gr.Textbox(label="Transcription")
+        persona_text = gr.Textbox(label="AI-Predicted Persona")
 
-    analyze_button.click(fn=analyze_speaker, inputs=audio_input, outputs=[transcribed_text, persona_output])
+    transcribe_button.click(transcribe_and_analyze, inputs=audio_input, outputs=[transcribed_text, persona_text])
 
 app.launch()