Upload v2 of URL classifier model (hybrid BERT + tabular)

Training Metrics:
- Eval Loss: 0.07472482323646545
- Eval F1 Macro: 0.9319480242737234
- Eval Accuracy: 0.9817232375979112

handler.py

@@ -0,0 +1,138 @@
+import re
+import os
+import json
+import torch
+import torch.nn as nn
+from urllib.parse import urlparse
+from transformers import AutoModel, AutoConfig, AutoTokenizer
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+PROFILE_SLUGS = re.compile(
+    r'/(profile|store|shop|freelancers?|biz|therapists?|counsellors?|'
+    r'restaurants?|menu|cottage|actors?|celebrants?|broker-finder|'
+    r'users?|usr|sellers?|vendors?|merchants?|dealers?|agents?|'
+    r'members?|str|book|booking|appointments?)(/|$)', re.IGNORECASE
+)
+
+NUM_TABULAR_FEATURES = 6
+NUMERIC_ID_IN_PATH = re.compile(r'/\d{3,}(/|$)')
+TABULAR_HIDDEN_SIZE = 128
+
+KNOWN_PLATFORMS_PATH = os.path.join(os.path.dirname(__file__), "known_platforms.json")
+with open(KNOWN_PLATFORMS_PATH) as _f:
+    KNOWN_PLATFORMS = set(json.load(_f))
+
+try:
+    import tldextract
+    _get_registered_domain = lambda url: tldextract.extract(url).registered_domain.lower()
+    _tld = lambda url: tldextract.extract(url).suffix.lower()
+except ImportError:
+    _get_registered_domain = lambda url: '.'.join(urlparse(url).netloc.lower().split('.')[-2:])
+    _tld = lambda url: urlparse(url).netloc.lower().split('.')[-1]
+
+_subdomain_dot_count = lambda url: max(0, urlparse(url).netloc.count('.') - 1)
+_path_depth = lambda url: len([s for s in urlparse(url).path.split('/') if s])
+
+extract_tabular_features = lambda url: [
+    1.0 if PROFILE_SLUGS.search(urlparse(url).path.lower()) else 0.0,
+    1.0 if _get_registered_domain(url) in KNOWN_PLATFORMS else 0.0,
+    min(_path_depth(url) / 10.0, 1.0),
+    min(_subdomain_dot_count(url) / 3.0, 1.0),
+    1.0 if NUMERIC_ID_IN_PATH.search(urlparse(url).path) else 0.0,
+    1.0 if _tld(url) == 'jp' else 0.0,
+]
+
+
+class UrlBertWithTabular(nn.Module):
+    def __init__(self, bert_model_name, num_labels, num_tabular_features=NUM_TABULAR_FEATURES):
+        super().__init__()
+        self.bert = AutoModel.from_pretrained(bert_model_name)
+        self.hidden_size = self.bert.config.hidden_size
+        self.num_labels = num_labels
+        self.num_tabular_features = num_tabular_features
+        self.tabular_proj = nn.Sequential(
+            nn.Linear(num_tabular_features, TABULAR_HIDDEN_SIZE),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+        )
+        self.classifier = nn.Linear(self.hidden_size + TABULAR_HIDDEN_SIZE, num_labels)
+
+    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, tabular_features=None, **kwargs):
+        bert_output = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
+        cls_output = bert_output.last_hidden_state[:, 0, :]
+        tabular_proj = self.tabular_proj(tabular_features.float())
+        combined = torch.cat([cls_output, tabular_proj], dim=1)
+        logits = self.classifier(combined)
+        return SequenceClassifierOutput(logits=logits)
+
+    @classmethod
+    def from_pretrained(cls, save_directory):
+        with open(os.path.join(save_directory, "tabular_config.json")) as f:
+            tabular_config = json.load(f)
+        bert_config = AutoConfig.from_pretrained(save_directory)
+        model = cls.__new__(cls)
+        nn.Module.__init__(model)
+        model.bert = AutoModel.from_config(bert_config)
+        model.hidden_size = bert_config.hidden_size
+        model.num_labels = tabular_config["num_labels"]
+        model.num_tabular_features = tabular_config["num_tabular_features"]
+        model.tabular_proj = nn.Sequential(
+            nn.Linear(model.num_tabular_features, TABULAR_HIDDEN_SIZE),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+        )
+        model.classifier = nn.Linear(model.hidden_size + TABULAR_HIDDEN_SIZE, model.num_labels)
+        safetensors_path = os.path.join(save_directory, "model.safetensors")
+        bin_path = os.path.join(save_directory, "pytorch_model.bin")
+        if os.path.exists(safetensors_path):
+            from safetensors.torch import load_file
+            state_dict = load_file(safetensors_path)
+        else:
+            state_dict = torch.load(bin_path, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict)
+        return model
+
+
+LABEL_MAP = {0: "official_website", 1: "platform"}
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model = UrlBertWithTabular.from_pretrained(path)
+        self.model.eval()
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+
+    def __call__(self, data):
+        inputs = data.get("inputs", data)
+        if isinstance(inputs, str):
+            inputs = [inputs]
+
+        encodings = self.tokenizer(
+            inputs, padding=True, truncation=True, max_length=128, return_tensors="pt"
+        ).to(self.device)
+
+        tabular = torch.tensor(
+            [extract_tabular_features(url) for url in inputs], dtype=torch.float32
+        ).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model(
+                input_ids=encodings["input_ids"],
+                attention_mask=encodings["attention_mask"],
+                tabular_features=tabular,
+            )
+
+        probs = torch.softmax(outputs.logits, dim=-1)
+        results = []
+        for i in range(len(inputs)):
+            scores = probs[i].tolist()
+            predictions = [
+                {"label": LABEL_MAP.get(j, f"LABEL_{j}"), "score": scores[j]}
+                for j in range(len(scores))
+            ]
+            predictions.sort(key=lambda x: x["score"], reverse=True)
+            results.append(predictions)
+
+        return results

known_platforms.json

@@ -0,0 +1,254 @@
+[
+  "alibaba.com",
+  "vinted.com",
+  "vinted.co.uk",
+  "jalan.net",
+  "instagram.com",
+  "linkedin.com",
+  "facebook.com",
+  "salla.sa",
+  "trustpilot.com",
+  "youtube.com",
+  "twitter.com",
+  "tiktok.com",
+  "twitch.tv",
+  "crunchbase.com",
+  "wa.me",
+  "zid.store",
+  "pinterest.com",
+  "whatsapp.com",
+  "ebay.co.uk",
+  "myshopify.com",
+  "etsy.com",
+  "discord.gg",
+  "wixsite.com",
+  "square.site",
+  "reddit.com",
+  "onlyfans.com",
+  "vinted.co.uk",
+  "yelp.com",
+  "stan.store",
+  "threads.net",
+  "bigcartel.com",
+  "t.me",
+  "snapchat.com",
+  "fresha.com",
+  "discord.com",
+  "planity.com",
+  "bbb.org",
+  "tip4serv.com",
+  "tumblr.com",
+  "sumupstore.com",
+  "booksy.com",
+  "depop.com",
+  "nextdoor.com",
+  "doctolib.fr",
+  "bookinbeautiful.com",
+  "bandcamp.com",
+  "spotlight.com",
+  "shopify.com",
+  "whatnot.com",
+  "squarespace.com",
+  "beacons.ai",
+  "fanvue.com",
+  "linktr.ee",
+  "company.site",
+  "webflow.io",
+  "wordpress.com",
+  "weebly.com",
+  "blogspot.com",
+  "ueniweb.com",
+  "canva.site",
+  "booking.com",
+  "godaddysites.com",
+  "nearcut.com",
+  "yell.com",
+  "substack.com",
+  "cargo.site",
+  "business.site",
+  "sell.app",
+  "goo.gl",
+  "g.co",
+  "g.page",
+  "google.com",
+  "hiboutik.com",
+  "inteletravel.uk",
+  "jimdofree.com",
+  "jimdo.com",
+  "just-eat.co.uk",
+  "ubereats.com",
+  "deliveroo.co.uk",
+  "gofundme.com",
+  "toasttab.com",
+  "holidaycottages.co.uk",
+  "imdb.com",
+  "upwork.com",
+  "fiverr.com",
+  "amazon.com",
+  "amazon.co.uk",
+  "amazon.co.jp",
+  "rakuten.co.jp",
+  "yahoo.co.jp",
+  "line.me",
+  "hirameki7.site",
+  "umin.ac.jp",
+  "guinot.com",
+  "motrio.fr",
+  "chiens-de-france.com",
+  "counselling-directory.org.uk",
+  "helloself.com",
+  "realpeople.co.uk",
+  "justmortgages.co.uk",
+  "linktree.com",
+  "shopifypreview.com",
+  "glossgenius.com",
+  "rekaz.io",
+  "systeme.io",
+  "direct.me",
+  "fato.me",
+  "notjusttravel.com",
+  "vivastreet.co.uk",
+  "mysellauth.com",
+  "garage-auto.info",
+
+  "alibaba.com",
+  "airbnb.com",
+  "airbnb.co.uk",
+  "airbnb.com.sg",
+  "airbnb.com.au",
+  "vinted.com",
+  "vinted.fr",
+  "vinted.de",
+  "herbalife.com",
+  "gonnaorder.com",
+  "keeq.io",
+  "soundbetter.com",
+  "10web-site.ai",
+  "github.io",
+  "sellsn.io",
+  "egift-store.com",
+  "tdc.ne.jp",
+
+  "ebay.com",
+  "ebay.de",
+  "ebay.fr",
+  "ebay.com.au",
+  "amazon.de",
+  "amazon.fr",
+  "amazon.es",
+  "amazon.it",
+  "amazon.in",
+  "amazon.com.au",
+  "rakuten.com",
+
+  "tripadvisor.com",
+  "tripadvisor.co.uk",
+  "tripadvisor.co.jp",
+  "zillow.com",
+  "realtor.com",
+  "rightmove.co.uk",
+  "zoopla.co.uk",
+  "indeed.com",
+  "glassdoor.com",
+  "glassdoor.co.uk",
+  "monster.com",
+  "freelancer.com",
+  "peopleperhour.com",
+  "toptal.com",
+  "99designs.com",
+
+  "x.com",
+  "bsky.app",
+  "mastodon.social",
+  "spotify.com",
+  "soundcloud.com",
+  "deezer.com",
+
+  "medium.com",
+  "patreon.com",
+  "ko-fi.com",
+  "buymeacoffee.com",
+  "gumroad.com",
+  "teachable.com",
+  "udemy.com",
+  "skillshare.com",
+  "coursera.org",
+  "eventbrite.com",
+  "eventbrite.co.uk",
+  "meetup.com",
+
+  "uber.com",
+  "lyft.com",
+  "doordash.com",
+  "grubhub.com",
+  "deliveroo.com",
+  "deliveroo.fr",
+  "deliveroo.com.au",
+  "just-eat.com",
+  "justeat.it",
+
+  "poshmark.com",
+  "mercari.com",
+  "rover.com",
+  "taskrabbit.com",
+  "thumbtack.com",
+  "temu.com",
+  "shein.com",
+  "wish.com",
+  "asos.com",
+  "zalando.com",
+  "zalando.co.uk",
+  "zalando.de",
+  "shopee.com",
+  "lazada.com",
+  "tokopedia.com",
+  "flipkart.com",
+  "olx.com",
+  "leboncoin.fr",
+  "allegro.pl",
+  "bol.com",
+  "cdiscount.com",
+  "fnac.com",
+  "otto.de",
+
+  "tabelog.com",
+  "hotpepper.jp",
+  "gnavi.co.jp",
+  "kakaku.com",
+  "zozo.jp",
+  "mercari.jp",
+  "minne.com",
+  "creema.jp",
+  "stores.jp",
+  "base.shop",
+  "booth.pm",
+  "note.com",
+  "ameblo.jp",
+  "fc2.com",
+
+  "wix.com",
+  "strikingly.com",
+  "site123.com",
+  "webnode.com",
+  "duda.co",
+  "format.com",
+  "cargocollective.com",
+  "contently.com",
+  "about.me",
+  "bio.link",
+  "lnk.bio",
+  "campsite.bio",
+  "carrd.co",
+  "typedream.com",
+  "super.so",
+
+  "behance.net",
+  "dribbble.com",
+  "deviantart.com",
+  "flickr.com",
+  "500px.com",
+  "vimeo.com",
+  "dailymotion.com",
+  "bookipi.com",
+  "odoo.com"
+]

requirements.txt

@@ -0,0 +1 @@
+tldextract