Updated Host

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.11-slim
+FROM python:3.12-slim
 
 WORKDIR /app
 
@@ -7,6 +7,8 @@ RUN pip install --no-cache-dir -r requirements.txt
 
 COPY . .
 
+ENV HF_HOME=/tmp/huggingface
+
 EXPOSE 7860
 
-CMD ["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "api:app", "--host", "0.0.0.0", "--port", "7860"]

api.py

@@ -5,13 +5,10 @@ Run with: uvicorn api:app --reload
 """
 
 import json
-import os
-from pathlib import Path
 from typing import Optional
 
 import email
 from email import policy as email_policy
-import numpy as np
 import torch
 from fastapi import FastAPI, HTTPException, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
@@ -24,11 +21,8 @@ from transformers import (
 
 # ── Config ────────────────────────────────────────────────────────────────────
 
-BASE_DIR = Path(__file__).parent
-MODELS_DIR = BASE_DIR / "models"
-
-ROBERTA_DIR = MODELS_DIR / "roberta_large_final"
-ELECTRA_DIR = MODELS_DIR / "electra_large_final"
+ROBERTA_REPO = "Dpedrinho01/trained_roberta_large"
+ELECTRA_REPO = "Dpedrinho01/trained_electra_large"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -45,7 +39,7 @@ app = FastAPI(
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["https://pedrinho-dev01.github.io/gone-phishing/"],
+    allow_origins=["*"],
     allow_methods=["*"],
     allow_headers=["*"],
 )
@@ -54,17 +48,20 @@ app.add_middleware(
 # ── Model loading ─────────────────────────────────────────────────────────────
 
 class ModelBundle:
-    def __init__(self, model_dir: Path, model_class, tokenizer_class=None):
-        self.model_dir = model_dir
-        self.tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
-        self.model = model_class.from_pretrained(str(model_dir))
+    def __init__(self, repo_id: str, model_class):
+        print(f"Loading {repo_id} …")
+        self.tokenizer = AutoTokenizer.from_pretrained(repo_id)
+        self.model = model_class.from_pretrained(repo_id)
         self.model.to(DEVICE)
         self.model.eval()
 
-        threshold_path = model_dir / "threshold_config.json"
+        # Load threshold from the repo's threshold_config.json
+        from huggingface_hub import hf_hub_download
+        threshold_path = hf_hub_download(repo_id=repo_id, filename="threshold_config.json")
         with open(threshold_path) as f:
             cfg = json.load(f)
         self.threshold: float = cfg["recommended_threshold"]
+        print(f"  ✓ {repo_id} loaded (threshold={self.threshold}, device={DEVICE})")
 
     @torch.no_grad()
     def predict_proba(self, text: str) -> float:
@@ -89,11 +86,9 @@ electra_bundle: Optional[ModelBundle] = None
 @app.on_event("startup")
 def load_models():
     global roberta_bundle, electra_bundle
-    print("Loading RoBERTa …")
-    roberta_bundle = ModelBundle(ROBERTA_DIR, RobertaForSequenceClassification)
-    print("Loading ELECTRA …")
-    electra_bundle = ModelBundle(ELECTRA_DIR, ElectraForSequenceClassification)
-    print(f"Models loaded on {DEVICE}.")
+    roberta_bundle = ModelBundle(ROBERTA_REPO, RobertaForSequenceClassification)
+    electra_bundle = ModelBundle(ELECTRA_REPO, ElectraForSequenceClassification)
+    print(f"All models ready on {DEVICE}.")
 
 
 # ── Schemas ───────────────────────────────────────────────────────────────────
@@ -217,7 +212,7 @@ def extract_text_from_eml(raw_bytes: bytes) -> str:
     if subject:
         parts.append(f"Subject: {subject}")
 
-    # From / To for extra signal
+    # From for extra signal
     from_addr = msg.get("from", "")
     if from_addr:
         parts.append(f"From: {from_addr}")
@@ -233,7 +228,6 @@ def extract_text_from_eml(raw_bytes: bytes) -> str:
                 # Fallback to HTML only if no plain text found
                 import html as html_lib
                 raw_html = part.get_content()
-                # Very light strip — remove tags
                 import re
                 text = re.sub(r"<[^>]+>", " ", raw_html)
                 text = html_lib.unescape(text)
@@ -267,5 +261,4 @@ async def predict_eml(file: UploadFile = File(...)):
     print(analyzed_text)
     print("=== [END EMAIL CONTENT] ===\n")
 
-    # Reuse the existing ensemble prediction logic
     return predict(PredictRequest(text=analyzed_text, model="ensemble"))

models/electra_large_final/.gitattributes

@@ -1 +0,0 @@
-model.safetensors filter=lfs diff=lfs merge=lfs -text

models/electra_large_final/config.json

@@ -1,34 +0,0 @@
-{
-  "add_cross_attention": false,
-  "architectures": [
-    "ElectraForSequenceClassification"
-  ],
-  "attention_probs_dropout_prob": 0.1,
-  "bos_token_id": null,
-  "classifier_dropout": null,
-  "dtype": "float32",
-  "embedding_size": 1024,
-  "eos_token_id": null,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
-  "hidden_size": 1024,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "is_decoder": false,
-  "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 512,
-  "model_type": "electra",
-  "num_attention_heads": 16,
-  "num_hidden_layers": 24,
-  "pad_token_id": 0,
-  "position_embedding_type": "absolute",
-  "summary_activation": "gelu",
-  "summary_last_dropout": 0.1,
-  "summary_type": "first",
-  "summary_use_proj": true,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.3.0",
-  "type_vocab_size": 2,
-  "use_cache": false,
-  "vocab_size": 30522
-}

models/electra_large_final/threshold_config.json

@@ -1,15 +0,0 @@
-{
-  "recommended_threshold": 0.35,
-  "standard_metrics": {
-    "accuracy": 0.9256,
-    "f1": 0.9051987767584098,
-    "precision": 0.9230769230769231,
-    "recall": 0.888
-  },
-  "custom_metrics": {
-    "accuracy": 0.9256,
-    "f1": 0.9055837563451776,
-    "precision": 0.9195876288659793,
-    "recall": 0.892
-  }
-}

models/electra_large_final/tokenizer_config.json

@@ -1,14 +0,0 @@
-{
-  "backend": "tokenizers",
-  "cls_token": "[CLS]",
-  "do_lower_case": true,
-  "is_local": false,
-  "mask_token": "[MASK]",
-  "model_max_length": 512,
-  "pad_token": "[PAD]",
-  "sep_token": "[SEP]",
-  "strip_accents": null,
-  "tokenize_chinese_chars": true,
-  "tokenizer_class": "BertTokenizer",
-  "unk_token": "[UNK]"
-}

models/electra_large_final/training_args.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3e251fe80c570139a5ddea6518864f1ccf76ef6536208c2d234507ba2c06c2b9
-size 4856

models/roberta_large_final/.gitattributes

@@ -1 +0,0 @@
-model.safetensors filter=lfs diff=lfs merge=lfs -text

models/roberta_large_final/config.json

@@ -1,28 +0,0 @@
-{
-  "add_cross_attention": false,
-  "architectures": [
-    "RobertaForSequenceClassification"
-  ],
-  "attention_probs_dropout_prob": 0.1,
-  "bos_token_id": 0,
-  "classifier_dropout": null,
-  "dtype": "float32",
-  "eos_token_id": 2,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
-  "hidden_size": 1024,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "is_decoder": false,
-  "layer_norm_eps": 1e-05,
-  "max_position_embeddings": 514,
-  "model_type": "roberta",
-  "num_attention_heads": 16,
-  "num_hidden_layers": 24,
-  "pad_token_id": 1,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.3.0",
-  "type_vocab_size": 1,
-  "use_cache": false,
-  "vocab_size": 50265
-}

models/roberta_large_final/threshold_config.json

@@ -1,15 +0,0 @@
-{
-  "recommended_threshold": 0.35,
-  "standard_metrics": {
-    "accuracy": 0.9352,
-    "f1": 0.916923076923077,
-    "precision": 0.9410526315789474,
-    "recall": 0.894
-  },
-  "custom_metrics": {
-    "accuracy": 0.9336,
-    "f1": 0.9150460593654043,
-    "precision": 0.9371069182389937,
-    "recall": 0.894
-  }
-}

models/roberta_large_final/tokenizer_config.json

@@ -1,16 +0,0 @@
-{
-  "add_prefix_space": false,
-  "backend": "tokenizers",
-  "bos_token": "<s>",
-  "cls_token": "<s>",
-  "eos_token": "</s>",
-  "errors": "replace",
-  "is_local": false,
-  "mask_token": "<mask>",
-  "model_max_length": 512,
-  "pad_token": "<pad>",
-  "sep_token": "</s>",
-  "tokenizer_class": "RobertaTokenizer",
-  "trim_offsets": true,
-  "unk_token": "<unk>"
-}

models/roberta_large_final/training_args.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cf7746da523087b4c98b10face3adad900b52a4c3ab325a7207442bec1e9eddb
-size 4856