import os
import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

def load_model():
    model_id = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
    access_token = os.getenv("hf_mistral_token")

    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True, token=access_token)

    quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype="float16"
    )

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        quantization_config=quant_config,
        device_map="auto",
        token=access_token
    )

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

def main():
    st.title("ChatGPT-Clone")

    if "generator" not in st.session_state:
        with st.spinner("Loading model..."):
            st.session_state.generator = load_model()

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

    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    if prompt := st.chat_input("Ask anything..."):
        st.session_state.messages.append({"role": "user", "content": prompt})

        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                result = st.session_state.generator(
                    prompt,
                    max_new_tokens=512,
                    temperature=0.7,
                    do_sample=True,
                )[0]["generated_text"]

                st.markdown(result)

        st.session_state.messages.append({"role": "assistant", "content": result})

if __name__ == "__main__":
    main()