back

app.py

@@ -1,12 +1,9 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import Dict, List, Optional
-import torch
-import torch.nn as nn
-import os
-import json
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from typing import Dict, List
 import uvicorn
-from transformers import AutoTokenizer, AutoConfig, AutoModel
+import torch
 
 app = FastAPI(
     title="CodeBERT Multitask Similarity API",
@@ -14,70 +11,17 @@ app = FastAPI(
     version="1.0.0"
 )
 
-# Define the MultitaskCodeSimilarityModel class
-class MultitaskCodeSimilarityModel(nn.Module):
-    def __init__(self, model_name, num_labels, tokenizer):
-        super().__init__()
-        self.tokenizer = tokenizer
-        self.config = AutoConfig.from_pretrained(model_name)
-        self.config.num_labels = num_labels
-        self.encoder = AutoModel.from_pretrained(model_name, config=self.config)
-        self.classifier = nn.Linear(self.config.hidden_size, num_labels)
-
-    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)
-        return logits
-
 # Load model and tokenizer
-try:
-    model_name = "HabibaElbehairy/codebert-multitask-similarity"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-    # Get the config to extract num_labels
-    config = AutoConfig.from_pretrained(model_name)
-    num_labels = getattr(config, "num_labels", 3)  # Default to 3 if not found
-    
-    # Initialize the custom model
-    model = MultitaskCodeSimilarityModel(model_name, num_labels=num_labels, tokenizer=tokenizer)
-    
-    # Load the weights - try different paths
-    try:
-        # 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 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"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 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"}
+model_name = "HabibaElbehairy/codebert-multitask-similarity"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 # Input schema definitions
 class SourceCode(BaseModel):
     class_name: str
     code: str
-    
 class TestCase(BaseModel):
     id: str
     test_fixture: str
@@ -85,74 +29,54 @@ class TestCase(BaseModel):
     code: str
     target_class: str
     target_method: List[str]
-    
 class SimilarityInput(BaseModel):
     pair_id: str
     source_code: SourceCode
     test_case_1: TestCase
     test_case_2: TestCase
-
-@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 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:
-        # Format input to match training format
-        combined_input = (
-            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}"
-        )
+    combined_input = (
+        f"[SOURCE CLASS]: {data.source_code.class_name}\n"
+        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)
+    inputs = tokenizer(combined_input, return_tensors="pt", padding=True, truncation=True).to(device)
 
-        # Get prediction from the model
-        with torch.no_grad():
-            # 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
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-        # Process results
-        probs = torch.softmax(logits, dim=-1)[0].cpu().tolist()
-        prediction = torch.argmax(logits, dim=-1).item()
-        
-        # Map prediction to class name
-        classification = label_to_class.get(prediction, "Unknown")
-        
-        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,
-            },
-            "probabilities": probs
-        }
-    
-    except Exception as e:
-        import traceback
-        print(traceback.format_exc())
-        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
+    probs = torch.softmax(outputs.logits, dim=-1)
+    score = torch.argmax(probs, dim=-1).item()
+
+    label_map = {
+        1: ("Duplicate", "Tests cover the same logic with similar inputs."),
+        2: ("Redundant", "Tests validate similar behavior but with slightly varied input."),
+        3: ("Distinct", "Tests verify opposite ends of battery status spectrum through different charge levels.")
+    }
+
+    label, explanation = label_map[score]
+
+    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": score,
+            "classification": label,
+            "explanation": explanation
+        },
+        "probabilities": probs[0].tolist()
+    }
 
 # This allows the app to run locally or in HF Spaces
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+
+    uvicorn.run(app, host="0.0.0.0", port=7860)