Refactor model loading and feature extraction; update logging and deployment info

app.py

@@ -1,36 +1,33 @@
 import os
-import time
 import logging
 import torch
 import torch.nn.functional as F
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModel
 from typing import List
 import uvicorn
 from datetime import datetime
-import requests
-import tempfile
-
-# Import our model and feature extraction function
-from model_definition import CodeSimilarityClassifier, extract_features
+from transformers import AutoTokenizer, AutoModel
 
-# Set up logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# Set up logging with more details
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler()]
+)
 logger = logging.getLogger(__name__)
 
-# System information - Updated with current values
-DEPLOYMENT_DATE = "2025-06-22 22:00:24"
+# System information
+DEPLOYMENT_DATE = "2025-06-22 22:05:09"
 DEPLOYED_BY = "habibaelbehairy"
 
 # Get device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
-# Your Hugging Face model repository
-REPO_ID = "HabibaElbehairy/CodeSimilarityClassifier"
-HF_TOKEN = os.getenv("HF_TOKEN", None)  # Optional: Set HF_TOKEN environment variable if the model is private
+# Your Hugging Face model repository - this is the key configuration
+MODEL_NAME = "HabibaElbehairy/CodeSimilarityClassifier"
 
 # Initialize FastAPI app
 app = FastAPI(
@@ -40,7 +37,7 @@ app = FastAPI(
     docs_url="/",
 )
 
-# Add CORS middleware to allow cross-origin requests
+# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -49,7 +46,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Define label to class mapping (0, 1, 2 for the new model)
+# Define label to class mapping
 label_to_class = {0: "Duplicate", 1: "Redundant", 2: "Distinct"}
 
 # Define input models for API
@@ -71,91 +68,186 @@ class SimilarityInput(BaseModel):
     test_case_1: TestCase
     test_case_2: TestCase
 
+# Define the model class
+class CodeSimilarityClassifier(torch.nn.Module):
+    def __init__(self, model_name="microsoft/codebert-base", num_labels=3):
+        super().__init__()
+        self.encoder = AutoModel.from_pretrained(model_name)
+        self.dropout = torch.nn.Dropout(0.1)
+
+        # Create a more powerful classification head
+        hidden_size = self.encoder.config.hidden_size
+
+        self.classifier = torch.nn.Sequential(
+            torch.nn.Linear(hidden_size, hidden_size),
+            torch.nn.LayerNorm(hidden_size),
+            torch.nn.GELU(),
+            torch.nn.Dropout(0.1),
+            torch.nn.Linear(hidden_size, 512),
+            torch.nn.LayerNorm(512),
+            torch.nn.GELU(),
+            torch.nn.Dropout(0.1),
+            torch.nn.Linear(512, num_labels)
+        )
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=True
+        )
+
+        pooled_output = outputs.pooler_output
+        logits = self.classifier(pooled_output)
+
+        return logits
+
+# Function to extract features from code
+def extract_features(source_code, test_code_1, test_code_2):
+    """Extract specific features to help the model identify similarities"""
+    import re
+    
+    # Extract test fixtures
+    fixture1 = re.search(r'TEST(?:_F)?\s*\(\s*(\w+)', test_code_1)
+    fixture1 = fixture1.group(1) if fixture1 else ""
+
+    fixture2 = re.search(r'TEST(?:_F)?\s*\(\s*(\w+)', test_code_2)
+    fixture2 = fixture2.group(1) if fixture2 else ""
+
+    # Extract test names
+    name1 = re.search(r'TEST(?:_F)?\s*\(\s*\w+\s*,\s*(\w+)', test_code_1)
+    name1 = name1.group(1) if name1 else ""
+
+    name2 = re.search(r'TEST(?:_F)?\s*\(\s*\w+\s*,\s*(\w+)', test_code_2)
+    name2 = name2.group(1) if name2 else ""
+
+    # Extract assertions
+    assertions1 = re.findall(r'(EXPECT_|ASSERT_)(\w+)', test_code_1)
+    assertions2 = re.findall(r'(EXPECT_|ASSERT_)(\w+)', test_code_2)
+
+    # Extract function/method calls
+    calls1 = re.findall(r'(\w+)\s*\(', test_code_1)
+    calls2 = re.findall(r'(\w+)\s*\(', test_code_2)
+
+    # Create explicit feature section
+    same_fixture = "SAME_FIXTURE" if fixture1 == fixture2 else "DIFFERENT_FIXTURE"
+    common_assertions = set([a[0] + a[1] for a in assertions1]).intersection(set([a[0] + a[1] for a in assertions2]))
+    common_calls = set(calls1).intersection(set(calls2))
+
+    features = (
+        f"METADATA: {same_fixture} | "
+        f"FIXTURE1: {fixture1} | FIXTURE2: {fixture2} | "
+        f"NAME1: {name1} | NAME2: {name2} | "
+        f"COMMON_ASSERTIONS: {len(common_assertions)} | "
+        f"COMMON_CALLS: {len(common_calls)} | "
+        f"ASSERTION_RATIO: {len(common_assertions)/(len(assertions1) + len(assertions2)) if assertions1 and assertions2 else 0}"
+    )
+
+    return features
+
 # Global variables for model and tokenizer
-model = None
 tokenizer = None
+model = None
 
-# Load model and tokenizer on startup
+# Load model and tokenizer on startup - COMPLETELY REWORKED
 @app.on_event("startup")
 async def startup_event():
-    global model, tokenizer
+    global tokenizer, model
+    
     try:
-        logger.info(f"Loading model and tokenizer from {REPO_ID}...")
+        logger.info("=== Starting model loading process ===")
         
-        # Create headers with token if available
-        headers = {}
-        if HF_TOKEN:
-            headers["Authorization"] = f"Bearer {HF_TOKEN}"
-        
-        # Load tokenizer directly from Hugging Face
+        # Step 1: Load the tokenizer first (simpler)
+        logger.info(f"Loading tokenizer from {MODEL_NAME}...")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
+            logger.info("✅ Base tokenizer loaded successfully")
+        except Exception as e:
+            logger.error(f"❌ Failed to load tokenizer: {str(e)}")
+            raise
+
+        # Step 2: Create model with base architecture
+        logger.info("Creating model architecture...")
         try:
-            tokenizer = AutoTokenizer.from_pretrained(REPO_ID, use_auth_token=HF_TOKEN)
-            logger.info("Tokenizer loaded successfully")
+            model = CodeSimilarityClassifier(model_name="microsoft/codebert-base")
+            logger.info("✅ Model architecture created successfully")
         except Exception as e:
-            logger.error(f"Error loading tokenizer: {str(e)}")
+            logger.error(f"❌ Failed to create model architecture: {str(e)}")
             raise
         
-        # Create model instance using our CodeSimilarityClassifier class
-        model = CodeSimilarityClassifier(model_name="microsoft/codebert-base")
-        logger.info("Model created successfully")
+        # Step 3: Load weights from locally saved file if available
+        logger.info("Looking for model weights...")
         
-        # Download the model weights file directly
-        try:
-            model_url = f"https://huggingface.co/{REPO_ID}/resolve/main/pytorch_model.bin"
-            logger.info(f"Downloading model weights from: {model_url}")
-            
-            with tempfile.NamedTemporaryFile(delete=False) as tmp:
-                # Make the request to download the file
-                response = requests.get(model_url, headers=headers, stream=True)
-                if not response.ok:
-                    logger.error(f"Failed to download model: {response.status_code} - {response.text}")
-                    raise Exception(f"Failed to download model: {response.status_code}")
+        # Try several possible locations
+        possible_paths = [
+            "pytorch_model.bin",
+            "/app/pytorch_model.bin",
+            "model/pytorch_model.bin",
+            os.path.join(os.getcwd(), "pytorch_model.bin")
+        ]
+        
+        model_loaded = False
+        for path in possible_paths:
+            if os.path.exists(path):
+                logger.info(f"Found model weights at {path}")
+                try:
+                    state_dict = torch.load(path, map_location=device)
+                    model.load_state_dict(state_dict)
+                    logger.info(f"✅ Successfully loaded weights from {path}")
+                    model_loaded = True
+                    break
+                except Exception as e:
+                    logger.warning(f"⚠️ Failed to load weights from {path}: {str(e)}")
+        
+        # If local loading fails, try fetching from HuggingFace
+        if not model_loaded:
+            logger.info("Attempting to fetch model directly from Hugging Face...")
+            try:
+                # Pull model from Hugging Face
+                temp_model = AutoModel.from_pretrained(MODEL_NAME)
+                logger.info("✅ Base model loaded from Hugging Face")
                 
-                # Save the file to a temporary location
-                for chunk in response.iter_content(chunk_size=8192):
-                    if chunk:
-                        tmp.write(chunk)
+                # Initialize our custom model with the base model's weights
+                model = CodeSimilarityClassifier(model_name="microsoft/codebert-base")
                 
-                tmp_name = tmp.name
-            
-            # Load the model weights
-            logger.info(f"Loading model weights from temporary file")
-            state_dict = torch.load(tmp_name, map_location=device)
-            model.load_state_dict(state_dict)
-            
-            # Clean up the temporary file
-            os.unlink(tmp_name)
-            logger.info("Model weights loaded successfully")
-        except Exception as e:
-            logger.error(f"Error downloading model weights: {str(e)}")
-            raise
+                # Only copy the encoder part (which is what we loaded from HF)
+                model.encoder = temp_model
+                logger.info("✅ Successfully initialized model with base encoder")
+                model_loaded = True
+            except Exception as e:
+                logger.error(f"❌ Failed to fetch from Hugging Face: {str(e)}")
+        
+        # Final check
+        if not model_loaded:
+            logger.error("❌ CRITICAL: Could not load model weights from any source")
+            raise RuntimeError("Failed to load model weights")
         
         # Move model to device and set to evaluation mode
         model.to(device)
         model.eval()
+        logger.info(f"✅ Model moved to {device} and set to evaluation mode")
+        logger.info("=== Model loading process complete ===")
         
-        logger.info("Model and tokenizer loaded successfully!")
     except Exception as e:
-        logger.error(f"Error loading model: {str(e)}")
+        logger.error(f"❌ CRITICAL ERROR in startup: {str(e)}")
         import traceback
         logger.error(traceback.format_exc())
-        model = None
-        tokenizer = None
+        if model is None:
+            logger.error("❌ Model object is None")
+        if tokenizer is None:
+            logger.error("❌ Tokenizer object is None")
+        raise
 
-@app.get("/health", tags=["Health"])
+@app.get("/health")
 async def health_check():
     """Health check endpoint that also returns deployment information"""
-    if model is None or tokenizer is None:
-        return {
-            "status": "error", 
-            "message": "Model or tokenizer not loaded",
-            "deployment_date": DEPLOYMENT_DATE,
-            "deployed_by": DEPLOYED_BY
-        }
+    status = "ok" if (model is not None and tokenizer is not None) else "error"
+    message = "Model and tokenizer loaded correctly" if status == "ok" else "Model or tokenizer not loaded"
     
     return {
-        "status": "ok", 
-        "model": REPO_ID, 
+        "status": status, 
+        "message": message,
+        "model": MODEL_NAME, 
         "device": str(device),
         "deployment_date": DEPLOYMENT_DATE,
         "deployed_by": DEPLOYED_BY,
@@ -166,11 +258,8 @@ async def health_check():
 async def predict(data: SimilarityInput):
     """
     Predict similarity class between two test cases for a given source class.
-    
-    Input schema follows the specified format with source_code, test_case_1, and test_case_2.
-    Uses feature extraction and the CodeSimilarityClassifier model for classification.
     """
-    if model is None:
+    if model is None or tokenizer is None:
         raise HTTPException(status_code=500, detail="Model not loaded correctly")
     
     try:
@@ -213,7 +302,7 @@ async def predict(data: SimilarityInput):
 
             # THIS IS WHERE THE MODEL IS CALLED
             with torch.no_grad():
-                # Our CodeSimilarityClassifier model returns logits directly
+                # Our custom model
                 logits = model(
                     input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"]
@@ -238,8 +327,7 @@ async def predict(data: SimilarityInput):
                 "score": api_score,
                 "classification": classification,
             },
-            "probabilities": probs,
-            "extracted_features": features  # Include extracted features for transparency
+            "probabilities": probs
         }
     
     except Exception as e:
@@ -249,8 +337,17 @@ async def predict(data: SimilarityInput):
         logger.error(error_trace)
         raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
 
-# Example endpoint
-@app.get("/example", response_model=SimilarityInput, tags=["Examples"])
+# Root and example routes remain the same
+@app.get("/")
+async def root():
+    return {
+        "message": "Test Similarity Analyzer API",
+        "documentation": "/docs",
+        "deployment_date": DEPLOYMENT_DATE,
+        "deployed_by": DEPLOYED_BY
+    }
+
+@app.get("/example", response_model=SimilarityInput)
 async def get_example():
     """Get an example input to test the API"""
     return SimilarityInput(
@@ -277,17 +374,5 @@ async def get_example():
         )
     )
 
-@app.get("/", tags=["Root"])
-async def root():
-    """
-    Root endpoint that provides API information and links to documentation.
-    """
-    return {
-        "message": "Test Similarity Analyzer API",
-        "documentation": "/docs",
-        "deployment_date": DEPLOYMENT_DATE,
-        "deployed_by": DEPLOYED_BY
-    }
-
 if __name__ == "__main__":
     uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)