Update app.py

app.py

@@ -17,6 +17,7 @@ model.eval()
 # Scoring function
 # -------------------------------
 def score_essay(essay: str):
+    # Empty essay → zero scores
     if not essay or not essay.strip():
         return {
             "Task Achievement": 0.0,
@@ -26,7 +27,7 @@ def score_essay(essay: str):
             "Overall": 0.0,
         }
 
-    # Tokenize input
+    # Tokenize input, truncate to 512 tokens
     inputs = tokenizer(
         essay,
         return_tensors="pt",
@@ -34,32 +35,30 @@ def score_essay(essay: str):
         max_length=512
     )
 
-    # Inference
+    # Model inference
     with torch.no_grad():
         outputs = model(**inputs)
 
-    # Raw logits from model
-    raw_scores = outputs.logits.squeeze().cpu().numpy()
+    # Raw logits from the model
+    preds = outputs.logits.squeeze().cpu().numpy()
 
     # -------------------------------
-    # IELTS calibration (IMPORTANT)
+    # CORRECT CALIBRATION
     # -------------------------------
-    # Convert logits → IELTS band scale
-    bands = 0.75 * raw_scores + 5.0
+    # Clip logits to 0-9 (IELTS band)
+    scores = np.clip(preds, 0, 9)
 
-    # Length penalty (IELTS-like)
-    word_count = len(essay.split())
+    # Round to nearest 0.5
+    rounded = np.round(scores * 2) / 2
 
+    # Optional: Penalize very short essays (like IELTS examiners)
+    word_count = len(essay.split())
     if word_count < 150:
-        bands -= 1.0
+        rounded -= 1.0
     elif word_count < 250:
-        bands -= 0.5
+        rounded -= 0.5
 
-    # Clamp to valid IELTS range
-    bands = np.clip(bands, 0.0, 9.0)
-
-    # Round to nearest 0.5
-    bands = np.round(bands * 2) / 2
+    rounded = np.clip(rounded, 0, 9)
 
     labels = [
         "Task Achievement",
@@ -69,7 +68,7 @@ def score_essay(essay: str):
         "Overall"
     ]
 
-    return {label: float(score) for label, score in zip(labels, bands)}
+    return {label: float(score) for label, score in zip(labels, rounded)}
 
 # -------------------------------
 # Gradio UI
@@ -86,7 +85,7 @@ with gr.Blocks() as demo:
         placeholder="Paste your IELTS essay here..."
     )
 
-    score_output = gr.Label(label="Estimated IELTS Band Scores")
+    score_output = gr.JSON(label="Estimated IELTS Band Scores")
 
     submit_btn = gr.Button("Score Essay")