import streamlit as st
from huggingface_hub import InferenceClient
from gtts import gTTS
import os

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def generate_response(message, system_message, max_tokens, temperature, top_p):
    messages = [{"role": "system", "content": system_message}]
    messages.append({"role": "user", "content": message})

    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token

    # Generate speech from the response text using gTTS
    tts = gTTS(text=response, lang='en')
    audio_file = "response.mp3"
    tts.save(audio_file)
    return response, audio_file

st.title("Hi (:")

message = st.text_input("Enter your message:")
system_message = st.text_input("System message:", value="You are a friendly Chatbot.")
max_tokens = st.slider("Max new tokens", 1, 2048, 512)
temperature = st.slider("Temperature", 0.1, 4.0, 0.7, step=0.1)
top_p = st.slider("Top-p (nucleus sampling)", 0.1, 1.0, 0.95, step=0.05)

if st.button("Generate Response"):
    response, audio_file = generate_response(message, system_message, max_tokens, temperature, top_p)
    st.write(response)

    # Add HTML to autoplay the audio
    audio_html = f"""
    <audio autoplay>
    <source src="data:audio/mpeg;base64,{open(audio_file, "rb").read().encode("base64").decode()}" type="audio/mp3">
    Your browser does not support the audio element.
    </audio>
    """
    st.markdown(audio_html, unsafe_allow_html=True)