distil bert was the best

app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import torch
-from transformers import BertTokenizer, BertForSequenceClassification
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
 from torch.nn.functional import softmax
 
 app = FastAPI(
@@ -14,8 +14,8 @@ app = FastAPI(
 class ContactDetector:
     def __init__(self):
         cache_dir = "/app/model_cache"
-        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', cache_dir=cache_dir)
-        self.model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2, cache_dir=cache_dir)
+        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased', cache_dir=cache_dir)
+        self.model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2, cache_dir=cache_dir)
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.model.to(self.device)
         self.model.eval()