Update chatbot.py

chatbot.py

@@ -2,7 +2,7 @@ import os
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationChain
 import langchain.globals
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer,  BitsAndBytesConfig
 import streamlit as st
 
 my_model_id = os.getenv('MODEL_REPO_ID', 'Default Value')
@@ -10,8 +10,12 @@ token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
 
 @st.cache_resource 
 def load_model():
+    quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    # bnb_4bit_compute_dtype=torch.bfloat16
+    )
     tokenizer = AutoTokenizer.from_pretrained(my_model_id)
-    model = AutoModelForCausalLM.from_pretrained(my_model_id)
+    model = AutoModelForCausalLM.from_pretrained(my_model_id, device_map="auto",quantization_config=quantization_config)
     
     return tokenizer,model