Update app.py

app.py

@@ -1,24 +1,33 @@
 import gradio as gr
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 import torch
-import librosa
+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 pretrained model and processor
+# Load Wav2Vec2 model and processor for speech recognition
 processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
-model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
+model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h").to(device)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# 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)
 
-# Transcription function
-def transcribe(audio_path):
-    if audio_path is None:
+# Function to transcribe audio to text
+def transcribe(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
 
-    # Load audio file and resample to 16kHz mono
-    audio_np, sample_rate = librosa.load(audio_path, sr=16000)
+    if isinstance(audio_np, np.ndarray) and audio_np.ndim > 1:
+        audio_np = np.mean(audio_np, axis=1)
 
-    # Process and transcribe
     input_values = processor(audio_np, sampling_rate=16000, return_tensors="pt").input_values.to(device)
     with torch.no_grad():
         logits = model(input_values).logits
@@ -26,16 +35,33 @@ def transcribe(audio_path):
     transcription = processor.decode(predicted_ids[0])
     return transcription.lower()
 
-# Gradio interface
+# 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)
+
+# 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
+
+# Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as app:
-    gr.Markdown("# Voice2PersonaAI")
-    gr.Markdown("Upload or record your voice, and this app will transcribe what you say.")
+    gr.Markdown("# Voice2Persona AI")
+    gr.Markdown("Upload or record your voice. We'll transcribe it and guess your fictional personality.")
 
     with gr.Row():
-        audio_input = gr.Audio(label="🎤 Record or Upload Your Voice", type="filepath")
-        output_text = gr.Textbox(label="📝 Transcribed Text")
+        audio_input = gr.Audio(source="microphone", type="numpy", label="🎤 Your Voice")
+        transcribed_text = gr.Textbox(label="📝 Transcription")
+        persona_output = gr.Textbox(label="🧠 Persona Analysis")
+
+    analyze_button = gr.Button("Analyze")
 
-    transcribe_button = gr.Button("Transcribe")
-    transcribe_button.click(fn=transcribe, inputs=audio_input, outputs=output_text)
+    analyze_button.click(fn=analyze_speaker, inputs=audio_input, outputs=[transcribed_text, persona_output])
 
 app.launch()