AbhishekBhavnani/Restaurant-Token-Classifier

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)