Update app.py

app.py

@@ -26,14 +26,24 @@ print('Retriever loaded successfully')
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
-base_model_name = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"  # Replace with your actual base model path
+base_model_name = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-model = AutoModel.from_pretrained(base_model_name)
-adapter_path = "mohamedalcafory/PubMed_Llama3.1_Based_model"  # Replace with your actual adapter path
+# Load in 4-bit with CPU offload using quantization_config
+# Removed load_in_4bit as it's redundant when using quantization_config
+model = AutoModelForCausalLM.from_pretrained(
+    base_model_name, 
+    device_map="auto", 
+    trust_remote_code=True,  # Required for some models
+    quantization_config=BitsAndBytesConfig(
+        load_in_4bit=True,  # Specify 4-bit quantization within BitsAndBytesConfig
+        load_in_8bit_fp32_cpu_offload=True  # Enable CPU offload
+    )
+)
+adapter_path = "mohamedalcafory/PubMed_Llama3.1_Based_model"
 model.load_adapter(adapter_path)
+
 # tokenizer = AutoTokenizer.from_pretrained("mohamedalcafory/PubMed_Llama3.1_Based_model")
 # model = AutoModelForCausalLM.from_pretrained("mohamedalcafory/PubMed_Llama3.1_Based_model")
 print(f'Model loaded successfully: {model}')