Update handler.py

handler.py

@@ -25,11 +25,11 @@ class ModelHandler:
 
     def load_model(self):
         # Load token as env var
-        model_id = "NiCETmtm/Llama3_kw_gen_new"
+        model_id = "NiCETmtm/llama3_torch"
         token = os.getenv("HF_API_TOKEN")
         # Load model & tokenizer
-        self.model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, from_tf=True)
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
+        self.model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, trust_remote_code=True, from_tf=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token, trust_remote_code=True)
 
     def predict(self, inputs):
         tokens = self.tokenizer(inputs, return_tensors="pt")