Rename inference.py to handler.py

inference.py → handler.py

@@ -3,19 +3,12 @@ import torch
 from huggingface_hub import HfApi
 import torch.nn.functional as F
 from peft import PeftModel
-import os
 
-hf_token = os.getenv("HF_TOKEN")
 
-
-if not hf_token:
-    raise ValueError("Hugging Face token not found in environment variables!")
-
-HfApi().set_access_token(hf_token)
 # Load model and tokenizer
 model_name = "munzirmuneer/phishing_url_gemma_pytorch" # Replace with your specific model
-model_name2 = "google/gemma-2b"
-tokenizer = AutoTokenizer.from_pretrained(model_name2, use_auth_token=True)
+#model_name2 = "google/gemma-2b"
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
 model = AutoModelForSequenceClassification.from_pretrained(model_name, use_auth_token=True)
 model = PeftModel.from_pretrained(model, model_name, use_auth_token=True)