Updated.

app.py

@@ -3,7 +3,7 @@ import pandas as pd
 import faiss
 import numpy as np
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from sentence_transformers import SentenceTransformer
 
 # Load retrieval corpus & FAISS index
@@ -12,14 +12,21 @@ index = faiss.read_index("faiss_index.bin")
 
 # Load embedding model
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-
-# Swap to BioMedLM 2.7B (CPU-friendly biomedical model)
 model_id = "stanford-crfm/BioMedLM"
 
+bnb_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_threshold=6.0,
+)
+
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-tokenizer.pad_token = tokenizer.eos_token  # fix padding issue
+tokenizer.pad_token = tokenizer.eos_token
 
-generation_model = AutoModelForCausalLM.from_pretrained(model_id)
+generation_model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    quantization_config=bnb_config,
+)
 
 def retrieve_top_k(query, k=5):
     query_embedding = embedding_model.encode([query]).astype("float32")