import streamlit as st
from llama_cpp import Llama

# Initialize the model
llm = Llama.from_pretrained(
    repo_id="Divyansh12/check",
    filename="unsloth.F16.gguf",  # Ensure this matches your model file name
    verbose=True,
    n_ctx=32768,
    n_threads=2,
    chat_format="chatml"
)

# Define the function to get responses from the model
def respond(message, history, system_message, max_tokens, temperature, top_p):
    messages = [{"role": "system", "content": system_message}]

    for user_message, assistant_message in history:
        if user_message:
            messages.append({"role": "user", "content": user_message})
        if assistant_message:
            messages.append({"role": "assistant", "content": assistant_message})

    messages.append({"role": "user", "content": message})

    response = ""
    # Stream the response from the model
    response_stream = llm.create_chat_completion(
        messages=messages,
        stream=True,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p
    )

    # Collect the response chunks
    for chunk in response_stream:
        if len(chunk['choices'][0]["delta"]) != 0 and "content" in chunk['choices'][0]["delta"]:
            response += chunk['choices'][0]["delta"]["content"]

    return response  # Return the full response

# Streamlit UI
st.title("Chatbot Application")
st.write("### Interact with the chatbot!")

# User input fields
system_message = st.text_input("System Message", value="You are a friendly Chatbot.")
user_message = st.text_area("Your Message:")
max_tokens = st.slider("Max New Tokens", min_value=1, max_value=2048, value=512)
temperature = st.slider("Temperature", min_value=0.1, max_value=4.0, value=0.7)
top_p = st.slider("Top-p (Nucleus Sampling)", min_value=0.1, max_value=1.0, value=0.95)

# Chat history
if 'history' not in st.session_state:
    st.session_state.history = []

if st.button("Send"):
    # Get the response from the model
    response = respond(user_message, st.session_state.history, system_message, max_tokens, temperature, top_p)

    # Add user message and model response to history
    st.session_state.history.append((user_message, response))

    # Display the chat history
    st.write("### Chat History")
    for user_msg, assistant_msg in st.session_state.history:
        st.write(f"**User:** {user_msg}")
        st.write(f"**Assistant:** {assistant_msg}")