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"  # 4-bit quantized model

    # Load Hugging Face token from environment variable
    access_token = os.getenv("hf_mistral_token")
    
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True, token=access_token)

    bnb_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=bnb_config,
        device_map="auto"
    )

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


def main():

    # st.set_page_config(page_title="ChatGPT Clone", page_icon="🤖")
    st.title("ChatGPT-Clone")

    # Load the generator model only once
    if "generator" not in st.session_state:
        with st.spinner("Loading model..."):
            st.session_state.generator = load_model()

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

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

    # Chat input
    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..."):
                # Call Mistral-7B API
                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()