try

app.py

@@ -4,13 +4,10 @@ from typing import Dict, List, Optional
 import torch
 import torch.nn as nn
 import os
-import re
+import json
 import uvicorn
 from transformers import AutoTokenizer, AutoConfig, AutoModel
 
-# Author: habibaelbehairy
-# Last updated: 2025-04-21 21:51:33
-
 app = FastAPI(
     title="CodeBERT Multitask Similarity API",
     description="Compare test case similarity using a fine-tuned CodeBERT model.",
@@ -27,45 +24,11 @@ class MultitaskCodeSimilarityModel(nn.Module):
         self.encoder = AutoModel.from_pretrained(model_name, config=self.config)
         self.classifier = nn.Linear(self.config.hidden_size, num_labels)
 
-        # For explanation generation
-        self.decoder_embedding = nn.Linear(self.config.hidden_size, self.config.hidden_size)
-        self.decoder = nn.GRU(
-            input_size=self.config.hidden_size,
-            hidden_size=self.config.hidden_size,
-            batch_first=True
-        )
-        self.explanation_head = nn.Linear(self.config.hidden_size, len(tokenizer))
-
-    def forward(self, input_ids, attention_mask, explanation_ids=None, explanation_mask=None):
+    def forward(self, input_ids, attention_mask):
         outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
         pooled = outputs.last_hidden_state[:, 0]
         logits = self.classifier(pooled)
-
-        explanation_logits = None
-        if explanation_ids is not None:
-            batch_size = input_ids.size(0)
-            seq_length = explanation_ids.size(1)
-
-            # Initialize decoder with pooled representation
-            decoder_input = self.decoder_embedding(pooled).unsqueeze(1).expand(-1, seq_length, -1)
-
-            # Run decoder
-            decoder_outputs, _ = self.decoder(decoder_input)
-
-            # Generate logits for each position
-            explanation_logits = self.explanation_head(decoder_outputs)
-
-        return logits, explanation_logits
-
-# Helper function to extract test fixture from test code
-def extract_test_fixture(test_code):
-    """Extract the test fixture name from the test code."""
-    match = re.search(r'TEST\(\s*(\w+)\s*,', test_code)
-    if not match:
-        match = re.search(r'TEST_F\(\s*(\w+)\s*,', test_code)
-    if match:
-        return match.group(1)
-    return None
+        return logits
 
 # Load model and tokenizer
 try:
@@ -79,35 +42,36 @@ try:
     # Initialize the custom model
     model = MultitaskCodeSimilarityModel(model_name, num_labels=num_labels, tokenizer=tokenizer)
     
-    # Try to load the model from different possible locations
+    # Load the weights - try different paths
     try:
-        # Try loading from standard path
+        # Try to load from hub path
         model.load_state_dict(torch.load(os.path.join(model_name, "pytorch_model.bin"), map_location="cpu"))
     except:
         try:
-            # Try local path
+            # Try local path relative to file
             model_weights_path = os.path.join(os.path.dirname(__file__), "pytorch_model.bin")
             if os.path.exists(model_weights_path):
                 model.load_state_dict(torch.load(model_weights_path, map_location="cpu"))
         except Exception as e:
-            print(f"Failed to load custom model: {e}")
-            # Fall back to standard model as last resort
+            print(f"Error loading weights: {e}")
+            # Try to use hub's model directly as fallback
             model = AutoModelForSequenceClassification.from_pretrained(model_name)
     
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
     model.eval()
+    print(f"Model loaded successfully and running on {device}")
+    
+    # Load label mapping or use default
+    # Default mapping - 1-indexed (adjust based on your trained model)
+    label_to_class = {1: "Duplicate", 2: "Redundant", 3: "Distinct"}
     
 except Exception as e:
-    print(f"Error loading model: {e}")
-    # Fall back to standard model
-    from transformers import AutoModelForSequenceClassification
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, from_tf=False)
-    model.to(device)
-    model.eval()
-
-# Default label mapping
-label_to_class = {0: "Duplicate", 1: "Redundant", 2: "Distinct"}
+    print(f"Error during model initialization: {e}")
+    # Create a dummy model for API documentation/testing
+    tokenizer = None
+    model = None
+    label_to_class = {1: "Duplicate", 2: "Redundant", 3: "Distinct"}
 
 # Input schema definitions
 class SourceCode(BaseModel):
@@ -131,106 +95,38 @@ class SimilarityInput(BaseModel):
 @app.get("/health")
 async def health_check():
     """Check if the API is up and running."""
-    return {
-        "status": "healthy", 
-        "model": "CodeBERT Multitask Similarity", 
-        "timestamp": "2025-04-21 21:51:33",
-        "user": "habibaelbehairy"
-    }
+    return {"status": "healthy", "model": "CodeBERT Multitask Similarity", "timestamp": "2025-04-21 20:08:44"}
 
 @app.post("/predict")
 async def predict(data: SimilarityInput):
     """
     Predict similarity class between two test cases for a given source class.
     """
+    if model is None:
+        raise HTTPException(status_code=500, detail="Model not loaded correctly")
+    
     try:
-        # Extract test related information
-        source_code = data.source_code.code
-        test_code_1 = data.test_case_1.code
-        test_code_2 = data.test_case_2.code
-        
-        # Extract class and method information
-        class_1 = data.test_case_1.target_class
-        class_2 = data.test_case_2.target_class
-        method_1 = data.test_case_1.target_method
-        method_2 = data.test_case_2.target_method
-        
-        # Apply heuristics like in your local testing
-        
-        # 1. If classes are different and not Unknown, tests are Distinct
-        if class_1 and class_2 and class_1 != "Unknown" and class_2 != "Unknown" and class_1 != class_2:
-            prediction = 2  # Distinct
-            classification = label_to_class[prediction]
-            return {
-                "pair_id": data.pair_id,
-                "test_case_1_name": data.test_case_1.name,
-                "test_case_2_name": data.test_case_2.name,
-                "similarity": {
-                    "score": prediction,
-                    "classification": classification,
-                    "explanation": ""
-                },
-                "probabilities": [0.1, 0.1, 0.8]  # Confidence towards Distinct
-            }
-        
-        # 2. If methods are completely different, tests are likely Distinct
-        if (method_1 and method_2 and set(method_1).isdisjoint(set(method_2)) and 
-            not (len(method_1) == 1 and len(method_2) == 1 and 
-                 ("test" in method_1[0].lower() or "test" in method_2[0].lower()))):
-            prediction = 2  # Distinct
-            classification = label_to_class[prediction]
-            return {
-                "pair_id": data.pair_id,
-                "test_case_1_name": data.test_case_1.name,
-                "test_case_2_name": data.test_case_2.name,
-                "similarity": {
-                    "score": prediction,
-                    "classification": classification,
-                    "explanation": ""
-                },
-                "probabilities": [0.1, 0.1, 0.8]  # Confidence towards Distinct
-            }
-        
-        # 3. Extract test fixtures to compare if available
-        fixture_1 = data.test_case_1.test_fixture
-        fixture_2 = data.test_case_2.test_fixture
-        
-        # If fixtures don't match, tests might be Distinct as well
-        if fixture_1 and fixture_2 and fixture_1 != fixture_2:
-            prediction = 2  # Distinct
-            classification = label_to_class[prediction]
-            return {
-                "pair_id": data.pair_id,
-                "test_case_1_name": data.test_case_1.name,
-                "test_case_2_name": data.test_case_2.name,
-                "similarity": {
-                    "score": prediction,
-                    "classification": classification,
-                    "explanation": ""
-                },
-                "probabilities": [0.1, 0.1, 0.8]  # Confidence towards Distinct
-            }
-        
-        # For all other cases, use the model for prediction
-        # Format input
+        # Format input to match training format
         combined_input = (
-            f"SOURCE CODE: {source_code}\n"
-            f"TEST 1: {test_code_1}\n"
-            f"TEST 2: {test_code_2}"
+            f"SOURCE CODE: {data.source_code.code}\n"
+            f"TEST 1: {data.test_case_1.code}\n"
+            f"TEST 2: {data.test_case_2.code}"
         )
 
         # Tokenize input
         inputs = tokenizer(combined_input, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
 
-        # Get prediction
+        # Get prediction from the model
         with torch.no_grad():
-            # Handle both custom and standard models
-            if isinstance(model, MultitaskCodeSimilarityModel):
-                logits, _ = model(
+            # Check if using custom model or fallback
+            if hasattr(model, 'forward'):
+                # Our custom model
+                logits = model(
                     input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"]
                 )
             else:
+                # Fallback to standard model
                 outputs = model(**inputs)
                 logits = outputs.logits
 
@@ -241,9 +137,6 @@ async def predict(data: SimilarityInput):
         # Map prediction to class name
         classification = label_to_class.get(prediction, "Unknown")
         
-        # Empty explanation as requested
-        explanation = ""
-
         return {
             "pair_id": data.pair_id,
             "test_case_1_name": data.test_case_1.name,
@@ -251,14 +144,12 @@ async def predict(data: SimilarityInput):
             "similarity": {
                 "score": prediction,
                 "classification": classification,
-                "explanation": explanation
             },
             "probabilities": probs
         }
     
     except Exception as e:
         import traceback
-        print(f"Error in prediction: {str(e)}")
         print(traceback.format_exc())
         raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")