Update app.py

app.py

@@ -3,21 +3,30 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 import numpy as np
 
+# -------------------------------
 # Load model and tokenizer
-model_name = "JacobLinCool/IELTS_essay_scoring_safetensors"
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Prediction function
-def score_essay(essay):
-    if not essay.strip():
-        return {"Task Achievement": 0,
-                "Coherence & Cohesion": 0,
-                "Vocabulary": 0,
-                "Grammar": 0,
-                "Overall": 0}
-
-    # Tokenize and truncate to max 512 tokens
+# -------------------------------
+MODEL_NAME = "JacobLinCool/IELTS_essay_scoring_safetensors"
+
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+model.eval()
+
+# -------------------------------
+# Scoring function
+# -------------------------------
+def score_essay(essay: str):
+    if not essay or not essay.strip():
+        return {
+            "Task Achievement": 0.0,
+            "Coherence & Cohesion": 0.0,
+            "Vocabulary": 0.0,
+            "Grammar": 0.0,
+            "Overall": 0.0,
+        }
+
+    # Tokenize input
     inputs = tokenizer(
         essay,
         return_tensors="pt",
@@ -25,29 +34,73 @@ def score_essay(essay):
         max_length=512
     )
 
-    # Run inference
+    # Inference
     with torch.no_grad():
         outputs = model(**inputs)
 
-    # Extract logits and normalize to 9
-    preds = outputs.logits.squeeze().numpy()
-    normalized = (preds / preds.max()) * 9
-    rounded = np.round(normalized * 2) / 2
+    # Raw logits from model
+    raw_scores = outputs.logits.squeeze().cpu().numpy()
+
+    # -------------------------------
+    # IELTS calibration (IMPORTANT)
+    # -------------------------------
+    # Convert logits → IELTS band scale
+    bands = 0.75 * raw_scores + 5.0
+
+    # Length penalty (IELTS-like)
+    word_count = len(essay.split())
+
+    if word_count < 150:
+        bands -= 1.0
+    elif word_count < 250:
+        bands -= 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
 
-    # Map labels
-    labels = ["Task Achievement", "Coherence & Cohesion", "Vocabulary", "Grammar", "Overall"]
-    return {label: float(score) for label, score in zip(labels, rounded)}
+    labels = [
+        "Task Achievement",
+        "Coherence & Cohesion",
+        "Vocabulary",
+        "Grammar",
+        "Overall"
+    ]
 
+    return {label: float(score) for label, score in zip(labels, bands)}
+
+# -------------------------------
 # Gradio UI
+# -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## Automated IELTS Essay Scoring")
-    gr.Markdown("Paste your essay below to get scores for all dimensions (Task, Coherence, Vocabulary, Grammar, Overall).")
+    gr.Markdown("## 📝 Automated IELTS Writing Scorer")
+    gr.Markdown(
+        "Paste your IELTS Task 2 essay below. "
+        "The system will estimate band scores for all four criteria and the overall band."
+    )
+
+    essay_input = gr.Textbox(
+        lines=12,
+        placeholder="Paste your IELTS essay here..."
+    )
 
-    essay_input = gr.Textbox(lines=10, placeholder="Paste your IELTS essay here...")
-    score_output = gr.Label()
+    score_output = gr.Label(label="Estimated IELTS Band Scores")
 
     submit_btn = gr.Button("Score Essay")
-    submit_btn.click(fn=score_essay, inputs=essay_input, outputs=score_output)
 
-# Launch app (for HF Spaces, leave default)
+    submit_btn.click(
+        fn=score_essay,
+        inputs=essay_input,
+        outputs=score_output
+    )
+
+    gr.Markdown(
+        "⚠️ **Note:** This is an AI-based estimator, not an official IELTS examiner score."
+    )
+
+# -------------------------------
+# Launch app
+# -------------------------------
 demo.launch()