app.py

import streamlit as st
from huggingface_hub import InferenceClient
from gtts import gTTS
import base64
from pydub import AudioSegment
import io
from pydub.playback import play

st.sidebar.width = 100
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
pre_prompt = "Te llamarás Xaman 3.0 (NO LO REPITAS) tu rol como IA conductual, tus principios son: existencialismo estóico y holístico."
pre_prompt_sent = False

def format_prompt(message, history):
    global pre_prompt_sent
    prompt = "<s>"

    if not pre_prompt_sent and all(f"[INST] {pre_prompt} [/INST]" not in user_prompt for user_prompt, _ in history):
        prompt += f"[INST] {pre_prompt} [/INST]"
        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, speed=1.3):
    tts = gTTS(text=text, lang='es')
    audio_file_path = 'output.mp3'
    tts.save(audio_file_path)
    sound = AudioSegment.from_mp3(audio_file_path)
    sound = sound.speedup(playback_speed=speed)
    sound.export(audio_file_path, format="mp3")

    return audio_file_path

def generate(user_input, history, temperature=None, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
    global pre_prompt_sent
    temperature = float(temperature) if temperature is not None else 0.9
    if temperature < 1e-2:
        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(user_input, history)
    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_path = text_to_speech(response, speed=1.3) 
    audio_file = open(audio_file_path, 'rb')
    audio_bytes = io.BytesIO(open(audio_file_path, 'rb').read())

    return response, audio_bytes

if "history" not in st.session_state:
    st.session_state.history = []

if st.session_state.history:
    user_input = st.session_state.history[-1][0]
    output, audio_bytes = generate(user_input, history=st.session_state.history)
    st.session_state.history[-1] = (user_input, output)
else:
    output, audio_bytes = "", ""

st.text_area("Respuesta de la IA", value=output, height=500, max_chars=600, key="output_text", disabled=True)
user_input = st.text_input(label="Ingresa tu mensaje", value="")

if user_input:
    output, audio_bytes = generate(user_input, history=st.session_state.history)
    st.session_state.history.append((user_input, output))

st.markdown(
    f"""
    <audio autoplay="autoplay" controls="controls" src="data:audio/mp3;base64,{base64.b64encode(audio_bytes).decode()}" type="audio/mp3" id="audio_player"></audio>
    """,
    unsafe_allow_html=True
)