Views
No views yet
pip install transformers torch datasets evaluate seqeval1from transformers import DistilBertForTokenClassification, DistilBertTokenizerFast
2import torch
3
4# Load quantized model
5model_path = "your_username/ner_distilbert_quantized"
6model = DistilBertForTokenClassification.from_pretrained(model_path)
7tokenizer = DistilBertTokenizerFast.from_pretrained(model_path)
8
9# Example text
10text = "Apple is looking at buying U.K. startup for $1 billion"
11
12# Tokenize and predict
13inputs = tokenizer(text.split(),
14 is_split_into_words=True,
15 return_tensors="pt",
16 padding="max_length",
17 truncation=True,
18 max_length=128)
19
20with torch.no_grad():
21 outputs = model(**inputs)
22
23# Process predictions
24predictions = torch.argmax(outputs.logits, dim=2)[0]
25word_ids = inputs.word_ids()
26
27entities = []
28current_entity = []
29current_label = None
30
31for idx, (word_id, pred_id) in enumerate(zip(word_ids, predictions)):
32 if word_id is None: # Skip special tokens
33 continue
34
35 pred_label = model.config.id2label[pred_id.item()]
36
37 if pred_label.startswith("B-"):
38 if current_entity:
39 entities.append((" ".join(current_entity), current_label))
40 current_entity = [text.split()[word_id]]
41 current_label = pred_label[2:]
42 elif pred_label.startswith("I-") and current_label == pred_label[2:]:
43 current_entity.append(text.split()[word_id])
44 else:
45 if current_entity:
46 entities.append((" ".join(current_entity), current_label))
47 current_entity = []
48 current_label = None
49
50print("Identified Entities:")
51for entity, label in entities:
52 print(f"{entity}: {label}")1quantized_model = torch.quantization.quantize_dynamic(
2 model,
3 {torch.nn.Linear},
4 dtype=torch.qint8
5).
├── model/ # Quantized model files
├── training_script.py # Fine-tuning code
├── inference_example.ipynb # Usage examples
├── requirements.txt # Dependencies
└── README.md # This documentation