import streamlit as st
from huggingface_hub import InferenceClient

client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")

def format_prompt(message, history):
    prompt = "<s>"
    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(prompt, history, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
    temperature = float(temperature)
    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(prompt, history)
    response = client.text_generation(formatted_prompt, **generate_kwargs)

    if isinstance(response, list) and len(response) > 0 and 'choices' in response[0]:
        return response[0]['choices'][0]['text']
    elif 'choices' in response:
        return response['choices'][0]['text']
    else:
        return response['text']
        
def main():
    st.title("Mistral 7B Chat Interface")

    # Sidebar for adjusting parameters
    st.sidebar.header("Adjust Parameters")
    temperature = st.sidebar.slider("Temperature", 0.0, 1.0, 0.9, step=0.05)
    max_new_tokens = st.sidebar.slider("Max new tokens", 0, 1048, 256, step=64)
    top_p = st.sidebar.slider("Top-p (nucleus sampling)", 0.0, 1.0, 0.90, step=0.05)
    repetition_penalty = st.sidebar.slider("Repetition penalty", 1.0, 2.0, 1.2, step=0.05)

    # Chat interface
    user_input = st.text_area("User Input:", "")
    history = []  # You need to manage the conversation history here

    if st.button("Send"):
        history.append(("User", user_input))
        bot_response = generate(user_input, history, temperature, max_new_tokens, top_p, repetition_penalty)
        history.append(("Bot", bot_response))

    st.text("Chat History:")
    for role, message in history:
        st.write(f"{role}: {message}")

if __name__ == "__main__":
    main()