Views
No views yet
bert-base-casedB-BRAND / I-BRAND – Sponsor/brand nameB-URL / I-URL – Sponsor link1import torch
2from transformers import BertTokenizerFast, BertForTokenClassification
3
4# Load model + tokenizer from Hugging Face Hub
5model_path = "Bhavya54/sponsor-ner-model"
6model = BertForTokenClassification.from_pretrained(model_path)
7tokenizer = BertTokenizerFast.from_pretrained(model_path)
8model.eval()
9
10# ID to label mapping
11id2label = {
12 0: "B-Brand",
13 1: "B-URL",
14 2: "I-Brand",
15 3: "I-URL",
16 4: "O"
17}
18
19# NER prediction function
20def predict_ner(text, model, tokenizer, id2label):
21 encoded = tokenizer(text, return_tensors="pt", truncation=True, return_offsets_mapping=True)
22 offset_mapping = encoded.pop("offset_mapping")
23
24 with torch.no_grad():
25 outputs = model(**encoded)
26
27 logits = outputs.logits
28 predictions = torch.argmax(logits, dim=2)
29
30 tokens = tokenizer.convert_ids_to_tokens(encoded["input_ids"][0])
31 labels = [id2label[p.item()] for p in predictions[0]]
32
33 entities = []
34 current_entity = ""
35 current_label = None
36
37 for token, label in zip(tokens, labels):
38 if token in ['[CLS]', '[SEP]', '[PAD]']:
39 continue
40
41 if label.startswith("B-"):
42 if current_entity:
43 entities.append((current_label, current_entity.strip()))
44 current_label = label[2:]
45 current_entity = token
46 elif label.startswith("I-") and current_label == label[2:]:
47 if token.startswith("##"):
48 current_entity += token[2:]
49 else:
50 if current_label == "URL":
51 current_entity += token
52 else:
53 current_entity += " " + token
54 else:
55 if current_entity:
56 entities.append((current_label, current_entity.strip()))
57 current_entity = ""
58 current_label = None
59
60 if current_entity:
61 entities.append((current_label, current_entity.strip()))
62
63 # Separate entities
64 brands = [ent.replace(" ##", "") for label, ent in entities if label == "Brand"]
65 urls = [ent.replace(" ", "").replace("##", "") for label, ent in entities if label == "URL"]
66
67 return brands, urls
68
69# Example usage
70text = "Graduation season hits different when you’re glowing and protected — ZEELOOL’s got you covered: https://zeelool.com//ytb?type=ZJTA855"
71brands, urls = predict_ner(text, model, tokenizer, id2label)
72
73print("brands =", brands)
74print("urls =", urls)