Anvilogic/URLGuardian

1import streamlit as st2from transformers import pipeline, AutoConfig, AutoModelForSequenceClassification, AutoTokenizer3 4@st.cache_resource5def load_classifier(model_path: str):6    id2label = {0: "Safe", 1: "Unsafe"}7    label2id = {"Safe": 0, "Unsafe": 1}8    config = AutoConfig.from_pretrained(model_path, id2label=id2label, label2id=label2id)9    model = AutoModelForSequenceClassification.from_pretrained(model_path, config=config)10    tokenizer = AutoTokenizer.from_pretrained(model_path)11    return pipeline("text-classification", model=model, tokenizer=tokenizer)12 13def defang_url(url: str) -> str:14    """15    Defangs the URL to prevent it from being clickable.16    This function replaces the protocol and dots.17    For example:18        https://example.com  -->  hxxps://example[.]com19    """20    # Replace the protocol21    if url.startswith("https://"):22        url = url.replace("https://", "hxxps://")23    elif url.startswith("http://"):24        url = url.replace("http://", "hxxp://")25    26    # Replace periods in the rest of the URL27    return url.replace(".", "[.]")28 29st.title("URL Typosquatting Detection with URLGuardian")30st.markdown(31    "This app uses the **URLGuardian** classifier developed by Anvilogic to detect potential suspicious URL. "32    "Enter a URL to assess!"33)34 35model_path = "./URLGuardian" 36classifier = load_classifier(model_path)37 38url = st.text_input("Enter the URL:", value="example.com")39 40if st.button("Check Safety of the url"):41    if url:42        result = classifier(url)[0]43        label = result["label"]44        score = result["score"]45        defanged_url = defang_url(url)46        if  label=='Safe':47            st.success(48                f"The URL '{defanged_url}' is considered safe with a confidence of {score * 100:.2f}%."49            )50        else:51            st.error(52                f"The URL '{defanged_url}' is considered suspicious with a confidence of {score * 100:.2f}%."53            )54        # Optionally, you can display the full result for debugging purposes:55        st.write("Full classification output:", result)56    else:57        st.error("Please enter a URL.")