edit app

app.py

@@ -2,11 +2,16 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 from typing import Dict
+import uvicorn
 import torch
 
-app = FastAPI()
+app = FastAPI(
+    title="CodeBERT Multitask Similarity API",
+    description="Compare test case similarity using a fine-tuned CodeBERT model.",
+    version="1.0.0"
+)
 
-# Load the pretrained model and tokenizer
+# Load model and tokenizer
 model_name = "HabibaElbehairy/codebert-multitask-similarity"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
@@ -31,10 +36,11 @@ class SimilarityInput(BaseModel):
     source_code: SourceCode
     test_case_1: TestCase
     test_case_2: TestCase
-
-@app.post("/predict")
+@app.post("/predict", tags=["Prediction"])
 async def predict(data: SimilarityInput):
-    # Prepare input string for model
+    """
+    Predict similarity class between two test cases for a given source class.
+    """
     combined_input = (
         f"[SOURCE CLASS]: {data.source_code.class_name}\n"
         f"[SOURCE CODE]: {data.source_code.code}\n"
@@ -42,17 +48,14 @@ async def predict(data: SimilarityInput):
         f"[TEST 2]: {data.test_case_2.code}"
     )
 
-    # Tokenize
     inputs = tokenizer(combined_input, return_tensors="pt", padding=True, truncation=True).to(device)
 
-    # Predict
     with torch.no_grad():
         outputs = model(**inputs)
 
     probs = torch.softmax(outputs.logits, dim=-1)
     score = torch.argmax(probs, dim=-1).item()
 
-    # Mapping index to label and explanation
     label_map = {
         0: ("Duplicate", "Tests cover the same logic with similar inputs."),
         1: ("Redundant", "Tests validate similar behavior but with slightly varied input."),
@@ -63,6 +66,8 @@ async def predict(data: SimilarityInput):
 
     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,
@@ -70,3 +75,8 @@ async def predict(data: SimilarityInput):
         },
         "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)