Update app.py

app.py

@@ -1,18 +1,18 @@
-from huggingface_hub import InferenceClient
-from audiorecorder import audiorecorder
 import streamlit as st
 import base64
-import torch
+import io
+from huggingface_hub import InferenceClient
+from gtts import gTTS
+from audiorecorder import audiorecorder
 import speech_recognition as sr
-from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC, Wav2Vec2CTCTokenizer, AutoTokenizer, AutoModelForSeq2SeqLM
+from pydub import AudioSegment
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
-model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
-tts_model_name = "facebook/mms-tts-spa"
-tokenizer = AutoTokenizer.from_pretrained(tts_model_name)
-tts_model = AutoModelForSeq2SeqLM.from_pretrained(tts_model_name)
+tokenizer = AutoTokenizer.from_pretrained("facebook/mms-fastspeech2-spa")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/mms-fastspeech2-spa")
 
-pre_prompt_text = "You are a behavioral AI, your answers should be brief, stoic and humanistic."
+pre_prompt_text = "Eres una IA conductual, tus respuestas deberán ser breves, estóicas y humanistas."
 
 if "history" not in st.session_state:
     st.session_state.history = []
@@ -30,18 +30,32 @@ def recognize_speech(audio_data, show_messages=True):
     try:
         audio_text = recognizer.recognize_google(audio, language="es-ES")
         if show_messages:
-            st.subheader("Recognized text:")
+            st.subheader("Texto Reconocido:")
             st.write(audio_text)
-            st.success("Completed.")
+            st.success("Reconocimiento de voz completado.")
     except sr.UnknownValueError:
-        st.warning("The audio could not be recognized. Did you try to record something?")
+        st.warning("No se pudo reconocer el audio. ¿Intentaste grabar algo?")
         audio_text = ""
     except sr.RequestError:
-        st.error("Talk to me to get started!")
+        st.error("Hablame para comenzar!")
         audio_text = ""
 
     return audio_text
 
+def format_prompt(message, history):
+    prompt = "<s>"
+
+    if not st.session_state.pre_prompt_sent:
+        prompt += f"[INST] {pre_prompt_text} [/INST]"
+        st.session_state.pre_prompt_sent = True
+
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+
 def generate(audio_text, history, temperature=None, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
     client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
@@ -63,33 +77,24 @@ def generate(audio_text, history, temperature=None, max_new_tokens=512, top_p=0.
 
     for response_token in stream:
         response += response_token.token.text
+    
+    response = ' '.join(response.split()).replace('</s>', '')
+    audio_file = text_to_speech(response, speed=1.3)
+    return response, audio_file
 
-    return response
-
-
-def format_prompt(message, history):
-    prompt = "<s>"
-
-    if not st.session_state.pre_prompt_sent:
-        prompt += f"[INST] {pre_prompt_text} [/INST]"
-        st.session_state.pre_prompt_sent = True
-
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
 
-def text_to_speech(text):
-    input_ids = tokenizer(text, return_tensors="pt").input_ids
+def text_to_speech(text, speed=1.3):
+    inputs = tokenizer(text, return_tensors="pt", padding="longest", truncation=True)
     with torch.no_grad():
-        logits = tts_model.generate(input_ids)
-    audio = processor.decode(logits[0], skip_special_tokens=True)
-    return audio
+        output = model(**inputs)
+    audio = output['logits']
+    audio_bytes = io.BytesIO()
+    torch.save(audio, audio_bytes)
+    audio_bytes.seek(0)
+    return base64.b64encode(audio_bytes.read()).decode()
 
 def main():
-    audio_data = st.audio_recorder("Push to Play", "Stop Recording...")
+    audio_data = audiorecorder("Presiona para hablar", "Deteniendo la grabación...")
 
     if not audio_data.empty():
         st.audio(audio_data.export().read(), format="audio/wav")
@@ -97,13 +102,11 @@ def main():
         audio_text = recognize_speech("audio.wav")
 
         if audio_text:
-            # Llama a la función generate para obtener la respuesta generada
-            generated_response = generate(audio_text, history=st.session_state.history)
-            output = text_to_speech(generated_response)
+            output, audio_file = generate(audio_text, history=st.session_state.history)  
 
-            if output is not None:
+            if audio_file is not None:
                 st.markdown(
-                    f"""<audio autoplay="autoplay" controls="controls" src="data:audio/mp3;base64,{base64.b64encode(output.encode()).decode()}" type="audio/mp3" id="audio_player"></audio>""",
+                    f"""<audio autoplay="autoplay" controls="controls" src="data:audio/mp3;base64,{base64.b64encode(audio_file.read()).decode()}" type="audio/mp3" id="audio_player"></audio>""",
                     unsafe_allow_html=True)
 
 if __name__ == "__main__":