reverting the model to the running one

app.py

@@ -1,13 +1,13 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import RobertaTokenizer, RobertaForSequenceClassification
 from torch.nn.functional import softmax
 import re
 
 app = FastAPI(
-    title="ContactShieldAI Detection API",
-    description="API for detecting contact information in text using ContactShieldAI",
+    title="Contact Information Detection API",
+    description="API for detecting contact information in text",
     version="1.0.0",
     docs_url="/"
 )
@@ -15,8 +15,8 @@ app = FastAPI(
 class ContactDetector:
     def __init__(self):
         cache_dir = "/app/model_cache"
-        self.tokenizer = AutoTokenizer.from_pretrained('xxparthparekhxx/ContactShieldAI', cache_dir=cache_dir)
-        self.model = AutoModelForSequenceClassification.from_pretrained('xxparthparekhxx/ContactShieldAI', cache_dir=cache_dir)
+        self.tokenizer = RobertaTokenizer.from_pretrained('roberta-base', cache_dir=cache_dir)
+        self.model = RobertaForSequenceClassification.from_pretrained('roberta-base', 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()
@@ -36,6 +36,7 @@ detector = ContactDetector()
 class TextInput(BaseModel):
     text: str
 
+
 def check_regex_patterns(text):
     patterns = [
         r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',  # Email
@@ -50,7 +51,8 @@ def check_regex_patterns(text):
             return True
     return False
 
-@app.post("/detect_contact", summary="Detect contact information in text using ContactShieldAI")
+
+@app.post("/detect_contact", summary="Detect contact information in text")
 async def detect_contact(input: TextInput):
     try:
         # First, check with regex patterns
@@ -62,14 +64,14 @@ async def detect_contact(input: TextInput):
                 "method": "regex"
             }
         
-        # If no regex patterns match, use the ContactShieldAI model
+        # If no regex patterns match, use the model
         probability = detector.detect_contact_info(input.text)
         is_contact = detector.is_contact_info(input.text)
         return {
             "text": input.text,
             "contact_probability": probability,
             "is_contact_info": is_contact,
-            "method": "ContactShieldAI"
+            "method": "model"
         }
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))