Update app.py

app.py

@@ -1,12 +1,15 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
 
 @st.cache_resource
 def load_classifier(model_path: str):
-    # Loads the URLGuardian classifier from the Hugging Face Hub.
-    return pipeline("text-classification", model=model_path)
+    id2label = {0: "Safe", 1: "Unsafe"}
+    label2id = {"Safe": 0, "Unsafe": 1}
+    config = AutoConfig.from_pretrained(model_path, id2label=id2label, label2id=label2id)
+    model = AutoModelForSequenceClassification.from_pretrained(model_path, config=config)
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    return pipeline("text-classification", model=model, tokenizer=tokenizer)
 
-# App Title and description
 st.title("URL Typosquatting Detection with URLGuardian")
 st.markdown(
     "This app uses the **URLGuardian** classifier by Anvilogic from the Hugging Face Hub to detect potential typosquatting. "
@@ -16,20 +19,14 @@ st.markdown(
 model_path = "./URLGuardian" 
 classifier = load_classifier(model_path)
 
-# URL inputs
 url = st.text_input("Enter the URL:", value="https://example.com")
 
-# Typosquatting detection on button click
 if st.button("Check Safety of the url"):
     if url:
-        # Run the classifier on the input URL
         result = classifier(url)[0]
         label = result["label"]
         score = result["score"]
-
-        # Display result based on the label
-        # Adjust the label checking logic based on the model's documentation.
-        if "safe" in label.lower():
+        if  label=='Safe':
             st.success(
                 f"The URL '{url}' is considered safe with a confidence of {score * 100:.2f}%."
             )