metadata
tags:
- token-classification
- distilbert
- named-entity-recognition
license: apache-2.0
datasets: restaurant-data
language:
- en
model-index:
- name: DistilBERT for Token Classification
results:
- task:
type: token-classification
name: Token Classification
metrics:
- type: f1
value: 0.92
name: F1 Score
DistilBERT for Token Classification
📝 Overview
This model is a fine-tuned version of DistilBERT for token classification on restaurant-related data. It is capable of recognizing various entities such as:
- Dishes
- Amenities
- Ratings
- Restaurant names
- Opening hours
- Locations
- Prices
- Cuisine types
🧠 Model Details
- Architecture: DistilBertForTokenClassification
- Hidden Size: 768
- FFN Inner Hidden Size: 3072
- Attention Heads: 12
- Layers: 6
- Max Position Embeddings: 512
- Dropout: 0.1
- Tokenizer: DistilBERT Tokenizer
- Transformers Version: 4.51.3
🏷️ Labels
The model supports the following 17 token classification labels:
| ID | Label |
|---|---|
| 0 | O |
| 1 | B-Amenity |
| 2 | I-Amenity |
| 3 | B-Dish |
| 4 | I-Dish |
| 5 | B-Rating |
| 6 | I-Rating |
| 7 | B-Restaurant_Name |
| 8 | I-Restaurant_Name |
| 9 | B-Hours |
| 10 | I-Hours |
| 11 | B-Location |
| 12 | I-Location |
| 13 | B-Price |
| 14 | I-Price |
| 15 | B-Cuisine |
| 16 | I-Cuisine |
🚀 Usage Example
from transformers import DistilBertTokenizer, DistilBertForTokenClassification
# Load tokenizer and model
tokenizer = DistilBertTokenizer.from_pretrained("AbhishekBhavnani/Restaurant-Token-Classifier")
model = DistilBertForTokenClassification.from_pretrained("AbhishekBhavnani/Restaurant-Token-Classifier")
# Example input
inputs = tokenizer("Find The Best Place To Eat Pizza in Ahmedabad", return_tensors="pt")
outputs = model(**inputs)
# Get predictions
predictions = outputs.logits.argmax(dim=-1)
print(predictions)