added bert again better outputs

app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import torch
-from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
+from transformers import BertTokenizer, BertForSequenceClassification
 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 = 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.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.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.model.to(self.device)
         self.model.eval()
@@ -27,7 +27,7 @@ class ContactDetector:
         probabilities = softmax(outputs.logits, dim=1)
         return probabilities[0][1].item()  # Probability of contact info
 
-    def is_contact_info(self, text, threshold=0.5):
+    def is_contact_info(self, text, threshold=0.45):
         return self.detect_contact_info(text) > threshold
 
 detector = ContactDetector()