Update app.py

app.py

@@ -12,16 +12,14 @@ from nltk.stem import PorterStemmer, WordNetLemmatizer
 from nltk.tokenize import word_tokenize
 from textblob import TextBlob
 
-# Download NLTK data (runs once)
+# Download NLTK data
 try:
     nltk.data.find('tokenizers/punkt')
 except LookupError:
     nltk.download('punkt')
     nltk.download('stopwords')
     nltk.download('wordnet')
-    nltk.download('averaged_perceptron_tagger')
 
-# Prefer MODEL_ID, fall back to HF_MODEL_ID, then default
 MODEL_ID = (
     os.environ.get("MODEL_ID")
     or os.environ.get("HF_MODEL_ID")
@@ -32,10 +30,10 @@ app = FastAPI(title="Phishing Text Classifier with Preprocessing", version="1.0.
 
 
 # ============================================================================
-# TEXT PREPROCESSING CLASS
+# TEXT PREPROCESSING CLASS (FOR ANALYSIS ONLY, NOT FOR MODEL INPUT)
 # ============================================================================
 class TextPreprocessor:
-    """Complete NLP preprocessing pipeline"""
+    """NLP preprocessing for analysis and feature extraction"""
     
     def __init__(self):
         self.stemmer = PorterStemmer()
@@ -44,15 +42,14 @@ class TextPreprocessor:
     
     def tokenize(self, text: str) -> List[str]:
         """Break text into tokens"""
-        text_lower = text.lower()
-        return word_tokenize(text_lower)
+        return word_tokenize(text.lower())
     
     def remove_stopwords(self, tokens: List[str]) -> List[str]:
         """Remove common stop words"""
         return [token for token in tokens if token.isalnum() and token not in self.stop_words]
     
     def stem(self, tokens: List[str]) -> List[str]:
-        """Reduce tokens to stems (Porter Stemmer)"""
+        """Reduce tokens to stems"""
         return [self.stemmer.stem(token) for token in tokens]
     
     def lemmatize(self, tokens: List[str]) -> List[str]:
@@ -60,15 +57,15 @@ class TextPreprocessor:
         return [self.lemmatizer.lemmatize(token) for token in tokens]
     
     def sentiment_analysis(self, text: str) -> Dict:
-        """Analyze sentiment polarity, subjectivity, and detect phishing indicators"""
+        """Analyze sentiment and phishing indicators"""
         blob = TextBlob(text)
-        polarity = blob.sentiment.polarity  # -1 (negative) to 1 (positive)
-        subjectivity = blob.sentiment.subjectivity  # 0 (objective) to 1 (subjective)
+        polarity = blob.sentiment.polarity
+        subjectivity = blob.sentiment.subjectivity
         
-        # Detect persuasive/emotional language (common in phishing)
         phishing_indicators = {
             "urgent_words": bool(re.search(r'\b(urgent|immediate|act now|verify|confirm|update|click|verify account)\b', text, re.IGNORECASE)),
             "threat_words": bool(re.search(r'\b(suspend|limited|expire|locked|disabled|restricted)\b', text, re.IGNORECASE)),
+            "suspicious_urls": bool(re.search(r'http\S+|www\S+', text)),
             "urgency_level": "HIGH" if re.search(r'\b(urgent|immediate|act now)\b', text, re.IGNORECASE) else "LOW"
         }
         
@@ -80,46 +77,20 @@ class TextPreprocessor:
             "phishing_indicators": phishing_indicators
         }
     
-    def clean_text(self, text: str) -> str:
-        """Clean URLs, special characters, extra spaces"""
-        # Remove URLs
-        text = re.sub(r'http\S+|www\S+', '', text)
-        # Remove email addresses
-        text = re.sub(r'\S+@\S+', '', text)
-        # Remove special characters but keep spaces
-        text = re.sub(r'[^a-zA-Z0-9\s]', '', text)
-        # Remove extra whitespace
-        text = re.sub(r'\s+', ' ', text).strip()
-        return text
-    
     def preprocess(self, text: str) -> Dict:
-        """Complete preprocessing pipeline"""
-        # Step 1: Clean
-        cleaned_text = self.clean_text(text)
-        
-        # Step 2: Tokenize
-        tokens = self.tokenize(cleaned_text)
-        
-        # Step 3: Remove stopwords
+        """Preprocessing for analysis (NOT for model)"""
+        tokens = self.tokenize(text)
         tokens_no_stop = self.remove_stopwords(tokens)
-        
-        # Step 4: Stem
         stemmed = self.stem(tokens_no_stop)
-        
-        # Step 5: Lemmatize
         lemmatized = self.lemmatize(tokens_no_stop)
-        
-        # Step 6: Sentiment analysis
         sentiment = self.sentiment_analysis(text)
         
         return {
             "original_text": text,
-            "cleaned_text": cleaned_text,
             "tokens": tokens,
             "tokens_without_stopwords": tokens_no_stop,
             "stemmed_tokens": stemmed,
             "lemmatized_tokens": lemmatized,
-            "processed_text": " ".join(lemmatized),  # Use lemmatized for model input
             "sentiment": sentiment,
             "token_count": len(tokens_no_stop)
         }
@@ -161,10 +132,7 @@ _NORM_LABELS_BY_IDX = None
 # HELPER FUNCTIONS
 # ============================================================================
 def _normalize_label_text_only(txt: str) -> str:
-    """
-    Normalize model label text to PHISH/LEGIT when possible.
-    If unfamiliar, return the uppercased original token.
-    """
+    """Normalize model label text"""
     t = (str(txt) if txt is not None else "").strip().upper()
     if t in ("PHISHING", "PHISH", "SPAM"):
         return "PHISH"
@@ -174,7 +142,7 @@ def _normalize_label_text_only(txt: str) -> str:
 
 
 def _load_model():
-    """Load model, tokenizer, and preprocessor on first use"""
+    """Load model, tokenizer, and preprocessor"""
     global _tokenizer, _model, _device, _NORM_LABELS_BY_IDX, _preprocessor
 
     if _tokenizer is None or _model is None:
@@ -184,44 +152,44 @@ def _load_model():
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
         _model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
         _model.to(_device)
-        _model.eval()  # important: disable dropout etc.
+        _model.eval()
         _preprocessor = TextPreprocessor()
 
-        # Warm-up (silent)
+        # Warm-up
         with torch.no_grad():
             _ = _model(
                 **_tokenizer(["warm up"], return_tensors="pt", padding=True, truncation=True, max_length=512)
                 .to(_device)
             ).logits
 
-        # Read and normalize model labels (by index)
+        # Read and normalize model labels
         id2label = getattr(_model.config, "id2label", {}) or {}
         num_labels = int(getattr(_model.config, "num_labels", 0) or 0)
         _NORM_LABELS_BY_IDX = [_normalize_label_text_only(id2label.get(i, f"LABEL_{i}")) for i in range(num_labels)]
         
-        print(f"Model loaded successfully. Number of labels: {num_labels}")
-        print(f"Label mapping: {id2label}")
+        print(f"Model loaded successfully")
+        print(f"ID2Label: {id2label}")
         print(f"Normalized labels: {_NORM_LABELS_BY_IDX}")
 
 
 def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List[Dict]:
     """
-    Predict and return strictly model-authoritative outputs with enhanced debugging.
+    Predict using ORIGINAL text (NO cleaning).
+    Preprocessing is for ANALYSIS only, not for model input.
     """
     _load_model()
     if not texts:
         return []
 
-    # Step 1: Preprocess texts if requested
+    # IMPORTANT: Use original text for model, NOT cleaned text!
+    model_inputs = texts
+    
+    # Get preprocessing info for analysis
     preprocessing_info = None
     if include_preprocessing:
         preprocessing_info = [_preprocessor.preprocess(text) for text in texts]
-        # Use lemmatized text for model input
-        model_inputs = [prep["processed_text"] for prep in preprocessing_info]
-    else:
-        model_inputs = texts
 
-    # Step 2: Tokenize batch for model
+    # Tokenize batch for model
     enc = _tokenizer(
         model_inputs,
         return_tensors="pt",
@@ -231,12 +199,11 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List
     )
     enc = {k: v.to(_device) for k, v in enc.items()}
 
-    # Step 3: Get predictions
+    # Get predictions
     with torch.no_grad():
         logits = _model(**enc).logits
-        probs = torch.softmax(logits, dim=-1)  # [batch, num_labels]
+        probs = torch.softmax(logits, dim=-1)
 
-    # Step 4: Build probability maps
     id2label = getattr(_model.config, "id2label", None) or {}
     labels_by_idx_raw = [id2label.get(i, f"LABEL_{i}") for i in range(probs.shape[-1])]
     labels_by_idx_norm = [_normalize_label_text_only(lbl) for lbl in labels_by_idx_raw]
@@ -249,7 +216,6 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List
         raw_label = labels_by_idx_raw[idx]
         norm_label = labels_by_idx_norm[idx]
 
-        # Build probability map keyed by normalized labels
         prob_map: Dict[str, float] = {}
         for j, lbl_norm in enumerate(labels_by_idx_norm):
             key = lbl_norm if lbl_norm in ("PHISH", "LEGIT") else f"CLASS_{j}"
@@ -260,13 +226,11 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List
             "raw_label": raw_label,
             "is_phish": True if norm_label == "PHISH" else False,
             "score": round(float(p[idx].item()), 4),
+            "confidence": round(float(p[idx].item()), 4),
             "probs": {k: round(v, 4) for k, v in prob_map.items()},
             "predicted_index": idx,
-            "all_logits": [round(float(logits[i][j].item()), 4) for j in range(logits.shape[1])],  # DEBUG
-            "raw_probs": [round(float(p[j].item()), 4) for j in range(len(p))],  # DEBUG
         }
         
-        # Add preprocessing info if requested
         if include_preprocessing and preprocessing_info:
             output["preprocessing"] = preprocessing_info[i]
         
@@ -281,26 +245,19 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List
 
 @app.get("/")
 def root():
-    """Root endpoint - shows API status"""
+    """Root endpoint"""
     _load_model()
     return {
         "status": "ok",
         "model": MODEL_ID,
         "device": _device,
-        "note": "Text preprocessing with stemming, lemmatization, stopword removal, and sentiment analysis enabled by default",
-        "endpoints": {
-            "/predict": "POST - Single text prediction",
-            "/predict-batch": "POST - Batch predictions",
-            "/evaluate": "POST - Evaluate with labeled samples",
-            "/debug/labels": "GET - View model label configuration",
-            "/debug/preprocessing": "POST - Debug preprocessing output only"
-        }
+        "note": "Model uses ORIGINAL text for predictions. Preprocessing is for analysis only.",
     }
 
 
 @app.get("/debug/labels")
 def debug_labels():
-    """Debug endpoint - view model label configuration"""
+    """View model configuration"""
     _load_model()
     return {
         "id2label": getattr(_model.config, "id2label", {}),
@@ -313,7 +270,7 @@ def debug_labels():
 
 @app.post("/debug/preprocessing")
 def debug_preprocessing(payload: PredictPayload):
-    """Debug endpoint - view preprocessing output only (no model prediction)"""
+    """Debug preprocessing output"""
     try:
         _load_model()
         preprocessing = _preprocessor.preprocess(payload.inputs)
@@ -327,7 +284,7 @@ def debug_preprocessing(payload: PredictPayload):
 
 @app.post("/predict")
 def predict(payload: PredictPayload):
-    """Single text prediction with optional preprocessing details"""
+    """Single prediction"""
     try:
         res = _predict_texts([payload.inputs], include_preprocessing=payload.include_preprocessing)
         return res[0]
@@ -337,7 +294,7 @@ def predict(payload: PredictPayload):
 
 @app.post("/predict-batch")
 def predict_batch(payload: BatchPredictPayload):
-    """Batch text predictions with optional preprocessing details"""
+    """Batch predictions"""
     try:
         return _predict_texts(payload.inputs, include_preprocessing=payload.include_preprocessing)
     except Exception as e:
@@ -346,10 +303,7 @@ def predict_batch(payload: BatchPredictPayload):
 
 @app.post("/evaluate")
 def evaluate(payload: EvalPayload):
-    """
-    Evaluate model on labeled samples.
-    Compares model predictions against provided ground truth labels.
-    """
+    """Evaluate on labeled samples"""
     try:
         texts = [s.text for s in payload.samples]
         gts = [(_normalize_label_text_only(s.label) if s.label is not None else None) for s in payload.samples]
@@ -383,6 +337,5 @@ def evaluate(payload: EvalPayload):
 
 
 if __name__ == "__main__":
-    # Run: uvicorn app:app --host 0.0.0.0 --port 8000 --reload
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)