Views
No views yet
Model Type: Transformer-based NER
Base Model: DistilBERT (distilbert-base-uncased)
Dataset: WNUT 17
Training Framework: PyTorch & Hugging Face Transformers
Training Epochs: 3
Batch Size: 16
Learning Rate: 2e-5
Optimizer: AdamW
Weight Decay: 0.01
Evaluation Strategy: Per epochperson
location
corporation
product
creative-work
group1#Loading the Model
2
3from transformers import DistilBertForTokenClassification, DistilBertTokenizerFast
4import torch
5
6model_name = "AventIQ-AI/distilbert-base-uncased_token_classification"
7
8def predict_entities(text, model, tokenizer):
9 """Predict Named Entities from the quantized model"""
10 inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
11
12 # Convert to FP32 if needed (for stability)
13 with torch.no_grad():
14 outputs = model(**inputs)
15 predictions = torch.argmax(outputs.logits.float(), dim=2) # Convert logits to float32
16
17 predicted_labels = [model.config.id2label[t.item()] for t in predictions[0]]
18 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
19
20 # Remove special tokens and align subwords
21 entities = []
22 current_entity = None
23
24 for token, label in zip(tokens, predicted_labels):
25 if token in [tokenizer.cls_token, tokenizer.sep_token, tokenizer.pad_token]:
26 continue
27
28 if token.startswith("##"): # Handle subwords
29 if current_entity:
30 current_entity["text"] += token[2:]
31 continue
32
33 if label == "O":
34 if current_entity:
35 entities.append(current_entity)
36 current_entity = None
37 else:
38 if label.startswith("B-"):
39 if current_entity:
40 entities.append(current_entity)
41 current_entity = {"text": token, "type": label[2:]}
42 elif label.startswith("I-") and current_entity:
43 current_entity["text"] += " " + token
44
45 if current_entity:
46 entities.append(current_entity)
47
48 return entities1test_sentence = ["Apple CEO Tim Cook announced the new iPhone 14 at their headquarters in Cupertino."]
2for sentence in test_sentences:
3 print(f"\nInput: {sentence}")
4 entities = predict_entities(sentence, model, tokenizer)
5 print("Detected entities:")
6 for entity in entities:
7 print(f"- {entity['text']} ({entity['type']})")
8 print("-" * 50)| Entity Type | Precision | Recall | F1 Score | Number of Entities |
|---|---|---|---|---|
| LOC (Location) | 91.46% | 92.07% | 91.76% | 3,000 |
| MISC (Miscellaneous) | 71.25% | 72.83% | 72.03% | 1,266 |
| ORG (Organization) | 89.83% | 93.02% | 91.40% | 3,524 |
| PER (Person) | 95.16% | 94.04% | 94.60% | 2,989 |
wnut_17 dataset was used, containing texts and their ner tags..
├── model/ # Contains the quantized model files
├── tokenizer_config/ # Tokenizer configuration and vocabulary files
├── model.safetensors/ # Quantized Model
├── README.md # Model documentation