Update app.py

app.py

@@ -3,11 +3,17 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
-base_model = "ybelkada/falcon-7b-sharded-bf16"
+base_model = "heliosbrahma/falcon-7b-sharded-bf16-finetuned-mental-health-conversational"
 adapter_model = "Sanjay002/falcon-7b-mental-health-finetuned"
 
-tokenizer = AutoTokenizer.from_pretrained(base_model)
-model = AutoModelForCausalLM.from_pretrained(base_model, device_map="auto", torch_dtype=torch.float16)
+tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    base_model,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True
+)
+
 model = PeftModel.from_pretrained(model, adapter_model)
 model.eval()