Update app.py

app.py

@@ -2,47 +2,43 @@ import os
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
-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)
 
-    # Proper way to load GPTQ models (auto-gptq must be installed)
+    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",
-        trust_remote_code=True,
         token=access_token
     )
 
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
     return pipe
 
-
 def main():
     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})
 
@@ -55,8 +51,9 @@ def main():
                     prompt,
                     max_new_tokens=512,
                     temperature=0.7,
-                    do_sample=True
+                    do_sample=True,
                 )[0]["generated_text"]
+
                 st.markdown(result)
 
         st.session_state.messages.append({"role": "assistant", "content": result})