updated app.py

app.py

@@ -15,7 +15,7 @@ from langchain_chroma import Chroma
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 from langchain_ollama import ChatOllama
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline,BitsAndBytesConfig
 from langchain.llms import HuggingFacePipeline
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
@@ -28,6 +28,7 @@ from random import randint
 import shutil
 from dotenv import load_dotenv
 from huggingface_hub import login
+import torch
 
 db_name = "vector_db"
 folder = "my-knowledge-base/"
@@ -37,9 +38,21 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN is None:
     raise ValueError("HF_TOKEN is not set. Check your .env file.")
 login(HF_TOKEN, add_to_git_credential=True)
-MODEL_NAME = "mistralai/Mistral-7B-4bit"
+
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,  # Set to 8-bit if needed
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto",use_auth_token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME, 
+    quantization_config=bnb_config, 
+    device_map="auto",
+    use_auth_token=HF_TOKEN
+)
 
 
 def process_files(files):