Update app.py

app.py

@@ -3,12 +3,20 @@ import time
 import torch
 
 from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-config = PeftConfig.from_pretrained("Ngadou/falcon-7b-scam-buster")
-model = AutoModelForCausalLM.from_pretrained("vilsonrodrigues/falcon-7b-instruct-sharded", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("vilsonrodrigues/falcon-7b-instruct-sharded")
-model = PeftModel.from_pretrained(model, "Ngadou/falcon-7b-scam-buster").to("cuda")
+peft_model_id = "Ngadou/falcon-7b-scam-buster"
+config = PeftConfig.from_pretrained(peft_model_id)
+
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, trust_remote_code=True, return_dict=True, load_in_8bit=True, device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+
+model = PeftModel.from_pretrained(model, peft_model_id).to("cuda")
+
+
+
+# Load the Lora model
+model = PeftModel.from_pretrained(model, peft_model_id)
 
 tokenizer.pad_token = tokenizer.eos_token