Update app.py

app.py

@@ -3,21 +3,26 @@ import os
 
 import torch
 from peft import PeftConfig, PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-base_model_name = "google/gemma-2b"
-#adapter_model_name = "samidh/cope-g2b-2c-hs-skr-s1.5.9-sx-sk-s5.d25"
-#adapter_model_name = "samidh/cope-g2b-2c-hs-skr-s1.5.9-sx-sk-s1.5.l1e4-e10-d25"
-#adapter_model_name = "samidh/cope-g2b-2c-hs-s1.f5.9.l5e5-e10-d25-r8"
-#adapter_model_name = "samidh/cope-g2b-2c-hs.s1.5.9-sx.s1.5.9o-hr.s5-sh.s5.l1e4-e10-d25-r8"
-#adapter_model_name = "samidh/cope-ap-g2b-2c-hs.s1.5.9-sx.s1.5.9o-vl.s1.5.9-hr.s5-sh.s5.l1e4-e5-d25-r8"
-#adapter_model_name = "samidh/cope-g2b-2c-hs.s1.5pcf.9.l5e5-e10-d25-r8"
-#adapter_model_name = "samidh/cope-g7bq-2c-hs.s1.5fpc.9-sx.s1.5.9o-VL.s1.5.9-HR.s5-SH.s5-l1e4-e5-d25-r8"
-adapter_model_name = "samidh/cope-g2b-2c-hs.s1.5fpc.9-sx.s1.5.9o-vl.s1.5.9-hr.s5-sh.s5.l5e5-e3-d25-r8"
-
-model = AutoModelForCausalLM.from_pretrained(base_model_name, token=os.environ['HF_TOKEN'])
+#base_model_name = "google/gemma-2b"
+base_model_name = "google/gemma-7b"
+#adapter_model_name = "samidh/cope-g2b-2c-hs.s1.5fpc.9-sx.s1.5.9o-vl.s1.5.9-hr.s5-sh.s5.l5e5-e3-d25-r8"
+adapter_model_name = "samidh/cope-g7bq-2c-hs.s1.5fpc.9-sx.s1.5.9o-VL.s1.5.9-HR.s5-SH.s5-l5e5-e3-d25-r8"
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    #bnb_4bit_quant_type="nf4",
+    #bnb_4bit_compute_dtype=torch.bfloat16,
+    #bnb_4bit_use_double_quant=True
+)
+
+model = AutoModelForCausalLM.from_pretrained(base_model_name,
+                                             token=os.environ['HF_TOKEN'],
+                                             quantization_config=bnb_config,
+                                             device_map="auto")
 model = PeftModel.from_pretrained(model, adapter_model_name, token=os.environ['HF_TOKEN'])
 model.merge_and_unload()
 
@@ -97,7 +102,7 @@ def predict(content, policy):
     input_text = PROMPT.format(policy=policy, content=content)
     input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
 
-    with torch.no_grad():
+    with torch.inference_mode():
         outputs = model(input_ids)
         logits = outputs.logits[:, -1, :]  # Get logits for the last token
         predicted_token_id = torch.argmax(logits, dim=-1).item()