Update app/main.py

app/main.py

@@ -1,24 +1,22 @@
 import os
-
-# Set a writable cache directory inside the Space container
-os.environ["TRANSFORMERS_CACHE"] = "/code/cache"
-
-
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import BertTokenizer, BertForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from sklearn.preprocessing import LabelEncoder
 import torch
 import numpy as np
 
+# Set cache
+os.environ["TRANSFORMERS_CACHE"] = "/code/cache"
+
 app = FastAPI()
 
-model = BertForSequenceClassification.from_pretrained(
+model = AutoModelForSequenceClassification.from_pretrained(
     "JaySenpai/bert-model",
     cache_dir="/code/cache",
-    use_safetensors=True  
+    use_safetensors=True
 )
-tokenizer = BertTokenizer.from_pretrained(
+tokenizer = AutoTokenizer.from_pretrained(
     "JaySenpai/bert-model",
     cache_dir="/code/cache"
 )
@@ -32,8 +30,7 @@ class TextInput(BaseModel):
 
 @app.post("/predict")
 async def predict(data: TextInput):
-    text = data.text  
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    inputs = tokenizer(data.text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)