Update app.py

app.py

@@ -1,20 +1,12 @@
 import streamlit as st
 import base64
 import io
-import torch
-import os
 from huggingface_hub import InferenceClient
-from audiorecorder import audiorecorder
+from gtts import gTTS
+import audiorecorder
 import speech_recognition as sr
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-import numpy as np
-from scipy.io.wavfile import write
-from pydub import AudioSegment
 
-model = T5ForConditionalGeneration.from_pretrained("facebook/mms-tts-spa")
-tokenizer = T5Tokenizer.from_pretrained("facebook/mms-tts-spa")
-
-pre_prompt_text = "Eres una IA conductual, tus respuestas deberán ser breves, estóicas y humanistas."
+pre_prompt_text = "You are a behavioral AI, your answers should be brief, stoic and humanistic."
 
 if "history" not in st.session_state:
     st.session_state.history = []
@@ -32,14 +24,14 @@ def recognize_speech(audio_data, show_messages=True):
     try:
         audio_text = recognizer.recognize_google(audio, language="es-ES")
         if show_messages:
-            st.subheader("Texto Reconocido:")
+            st.subheader("Recognized text:")
             st.write(audio_text)
-            st.success("Reconocimiento de voz completado.")
+            st.success("Voice Recognized.")
     except sr.UnknownValueError:
-        st.warning("No se pudo reconocer el audio. ¿Intentaste grabar algo?")
+        st.warning("The audio could not be recognized. Did you try to record something?")
         audio_text = ""
     except sr.RequestError:
-        st.error("Hablame para comenzar!")
+        st.error("Push/Talk to start!")
         audio_text = ""
 
     return audio_text
@@ -58,16 +50,41 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(audio_text, history):
+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")
+
+    temperature = float(temperature) if temperature is not None else 0.9
+    temperature = max(temperature, 1e-2)
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42)
+
     formatted_prompt = format_prompt(audio_text, history)
-    inputs = tokenizer(formatted_prompt, return_tensors="pt", padding=True, truncation=True)
-    with torch.no_grad():
-        output = model.generate(**inputs)
-    audio = output['audio']
-    return audio
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
+    response = ""
+
+    for response_token in stream:
+        response += response_token.token.text
+    
+    response = ' '.join(response.split()).replace('</s>', '')
+    audio_file = text_to_speech(response)
+    return response, audio_file
+
+def text_to_speech(text):
+    tts = gTTS(text=text, lang='es')
+    audio_fp = io.BytesIO()
+    tts.write_to_fp(audio_fp)
+    audio_fp.seek(0)
+    return audio_fp
 
 def main():
-    audio_data = audiorecorder("Presiona para hablar", "Deteniendo la grabación...")
+    audio_data = audiorecorder.audiorecorder("Push to Talk", "Stop Recording...")
 
     if not audio_data.empty():
         st.audio(audio_data.export().read(), format="audio/wav")
@@ -75,24 +92,12 @@ def main():
         audio_text = recognize_speech("audio.wav")
 
         if audio_text:
-            audio_file = generate(audio_text, history=st.session_state.history)
+            output, audio_file = generate(audio_text, history=st.session_state.history)  
 
             if audio_file is not None:
-                # Guardar el archivo WAV
-                write("output.wav", model.config.sampling_rate, audio_file)
-
-                # Convertir el archivo WAV a MP3 utilizando pydub
-                audio = AudioSegment.from_wav("output.wav")
-                audio.export("output.mp3", format="mp3")
-
-                # Leer el archivo MP3 y mostrarlo en Streamlit
-                with open("output.mp3", "rb") as file:
-                    audio_bytes = file.read()
-                st.audio(audio_bytes, format="audio/mp3")
-
-                # Eliminar archivos temporales (opcional)
-                os.remove("output.wav")
-                os.remove("output.mp3")
+                st.markdown(
+                    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__":
-    main()
+    main()