import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Initialize model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3.5-mini-instruct", 
    device_map="cpu",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3.5-mini-instruct")

# Create pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)

# Generation arguments
generation_args = {
    "max_new_tokens": 500,
    "return_full_text": False,
    "temperature": 0.0,
    "do_sample": False,
}

# Chat function
def chat(message, history, system_prompt):
    # Prepare messages
    messages = [
        {"role": "system", "content": system_prompt},
    ]
    
    # Add history to messages
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    # Generate response
    output = pipe(messages, **generation_args)
    response = output[0]['generated_text']
    
    return response

# Streamlit App Layout
st.title("Testing new Phi 3.5 model")
st.sidebar.title("Settings")
system_prompt = st.sidebar.text_input("System Prompt", value="You are a helpful AI assistant.")

# Initialize chat history
if "chat_history" not in st.session_state:
    st.session_state["chat_history"] = []

# Input text box
user_input = st.text_input("You:", key="input")

# Chat history display
if st.session_state["chat_history"]:
    for user_msg, bot_msg in st.session_state["chat_history"]:
        st.write(f"**You:** {user_msg}")
        st.write(f"**Bot:** {bot_msg}")

# On submit
if user_input:
    bot_response = chat(user_input, st.session_state["chat_history"], system_prompt)
    st.session_state["chat_history"].append((user_input, bot_response))
    #st.experimental_rerun()

# Clear chat history button
if st.button("Clear Chat"):
    st.session_state["chat_history"] = []
    #st.experimental_rerun()