updating this in prep to ship new single model with different hyperparameters

api/predict.py

@@ -6,42 +6,34 @@ from huggingface_hub import hf_hub_download
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer = None
-models = None
+model = None
 
 def load_resources():
-    global tokenizer, models
+    global tokenizer, model
     
-    if tokenizer is not None and models is not None:
+    if tokenizer is not None and model is not None:
         return
     
-    print("loading models...")
+    print("loading model...")
     
     tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
     
-    num_classes = 2
-    dropout = 0.4
-    
-    models = []
-    for i in range(1, 6):
-        model_filename = f"ensemble_model_{i}.pth"
-        
-        print(f"downloading {model_filename}...")
-        model_path = hf_hub_download(
-            repo_id="codingcoolfun9ed/sentinelcheck-models",
-            filename=model_filename
-        )
-        
-        model = DistilBertForSequenceClassification.from_pretrained(
-            'distilbert-base-uncased',
-            num_labels=num_classes,
-            dropout=dropout
-        )
-        model.load_state_dict(torch.load(model_path, map_location=device))
-        model = model.to(device)
-        model.eval()
-        models.append(model)
-    
-    print("models loaded")
+    print("downloading model_2.pth...")
+    modelPath = hf_hub_download(
+        repo_id="codingcoolfun9ed/sentinelcheck-models",
+        filename="model_2.pth"
+    )
+    
+    model = DistilBertForSequenceClassification.from_pretrained(
+        'distilbert-base-uncased',
+        num_labels=2,
+        dropout=0.4
+    )
+    model.load_state_dict(torch.load(modelPath, map_location=device))
+    model = model.to(device)
+    model.eval()
+    
+    print("model loaded")
 
 def cleanText(text):
     if not text:
@@ -88,33 +80,33 @@ def predict_review(text):
         return_tensors='pt'
     )
     
-    input_ids = encoding['input_ids'].to(device)
-    attention_mask = encoding['attention_mask'].to(device)
+    inputIds = encoding['input_ids'].to(device)
+    attentionMask = encoding['attention_mask'].to(device)
     
-    allOutputs = []
     with torch.no_grad():
-        for model in models:
-            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-            probs = torch.softmax(outputs.logits, dim=1)
-            allOutputs.append(probs.cpu().numpy())
+        outputs = model(input_ids=inputIds, attention_mask=attentionMask)
+        probs = torch.softmax(outputs.logits, dim=1).cpu().numpy()[0]
     
-    avgProbs = np.mean(allOutputs, axis=0)[0]
-    fakeProb = avgProbs[1]
-    realProb = avgProbs[0]
+    fakeProb = probs[1]
+    realProb = probs[0]
     
-    isFake = fakeProb > 0.75
     confidence = max(fakeProb, realProb)
-    prediction = "fake" if isFake else "real"
     
     if confidence < 0.75:
         prediction = "uncertain"
+        isFake = None
+    else:
+        isFake = fakeProb > realProb
+        prediction = "fake" if isFake else "real"
     
     lengthCat = getLengthCategory(cleaned)
     
     return {
         "prediction": prediction,
         "confidence": float(confidence),
-        "is_fake": bool(isFake),
+        "is_fake": isFake,
         "length_category": lengthCat,
-        "token_count": len(cleaned.split())
+        "token_count": len(cleaned.split()),
+        "fake_probability": float(fakeProb),
+        "real_probability": float(realProb)
     }