| """One essay in, the coach's reading out: six scores, argument parts with effectiveness, and plain-words notes. |
| |
| coach = Coach("/data/runs/scores-v1", "/data/runs/parts-v1") |
| coach.read(text) # dict |
| """ |
| import json |
| import os |
|
|
| import numpy as np |
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| from coach.data import TRAITS, TYPES, with_breaks |
| from coach.decode import decode |
| from coach.model import PartsModel, ScoreModel |
|
|
| TRAIT_WORDS = { |
| "cohesion": "how the ideas link up and flow from one to the next", |
| "syntax": "sentence structure and variety", |
| "vocabulary": "the range of words used", |
| "phraseology": "natural phrasing and word combinations", |
| "grammar": "grammatical accuracy", |
| "conventions": "spelling, punctuation and capitalisation", |
| } |
|
|
|
|
| class Coach: |
| def __init__(self, scores_dir, parts_dir, device=None): |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| sc = json.load(open(os.path.join(scores_dir, "config.json"))) |
| self.score_tok = AutoTokenizer.from_pretrained(os.path.join(scores_dir, "tokenizer")) |
| self.score_net = ScoreModel(sc["backbone"], n_tokens=len(self.score_tok)).to(self.device).eval() |
| self.score_net.load_state_dict(load_file(os.path.join(scores_dir, "model.safetensors"))) |
| self.score_len = sc["max_len"] |
| pc = json.load(open(os.path.join(parts_dir, "config.json"))) |
| self.parts_tok = AutoTokenizer.from_pretrained(os.path.join(parts_dir, "tokenizer")) |
| self.parts_net = PartsModel(pc["backbone"], n_tokens=len(self.parts_tok)).to(self.device).eval() |
| self.parts_net.load_state_dict(load_file(os.path.join(parts_dir, "model.safetensors"))) |
| self.parts_len = pc["max_len"] |
|
|
| @torch.no_grad() |
| def scores(self, text): |
| enc = self.score_tok(with_breaks(text), truncation=True, max_length=self.score_len, return_tensors="pt").to(self.device) |
| out = self.score_net(enc["input_ids"], enc["attention_mask"])[0].float().cpu().numpy() |
| return {t: float(np.clip(round(v * 2) / 2, 1.0, 5.0)) for t, v in zip(TRAITS, out)}, {t: float(v) for t, v in zip(TRAITS, out)} |
|
|
| @torch.no_grad() |
| def parts(self, text): |
| enc = self.parts_tok(text, truncation=True, max_length=self.parts_len, return_offsets_mapping=True, return_tensors="pt") |
| offs = [None if (s == e == 0) else (int(s), int(e)) for s, e in enc.pop("offset_mapping")[0].tolist()] |
| offs[0] = None |
| enc = enc.to(self.device) |
| lb, le = self.parts_net(enc["input_ids"], enc["attention_mask"]) |
| return decode(lb[0].float().cpu().numpy(), le[0].float().cpu().numpy(), offs, text) |
|
|
| def read(self, text): |
| text = text.replace("\r\n", "\n").strip() |
| snapped, raw = self.scores(text) |
| parts = self.parts(text) |
| return {"scores": snapped, "scores_raw": raw, "parts": parts, "notes": notes(snapped, parts, text)} |
|
|
|
|
| def notes(scores, parts, text): |
| """Plain-words feedback from the numbers: the two weakest skills, and what is missing or weak in the argument.""" |
| out = [] |
| weakest = sorted(scores, key=scores.get)[:2] |
| for t in weakest: |
| out.append({"kind": "skill", "trait": t, "score": scores[t], "text": f"Weakest area: {TRAIT_WORDS[t]} ({scores[t]:.1f} of 5)."}) |
| have = {p["type"] for p in parts} |
| if "Position" not in have: |
| out.append({"kind": "structure", "text": "No clear position: say early on which side you take, in one sentence."}) |
| if "Lead" not in have: |
| out.append({"kind": "structure", "text": "No lead: the first lines should pull the reader in and point at the question."}) |
| if "Counterclaim" not in have: |
| out.append({"kind": "structure", "text": "No counterclaim: name the strongest view against yours, then answer it. Examiners look for this."}) |
| elif "Rebuttal" not in have: |
| out.append({"kind": "structure", "text": "You raise the other side but never answer it: add a rebuttal after the counterclaim."}) |
| if "Concluding Statement" not in have: |
| out.append({"kind": "structure", "text": "No conclusion: finish by restating your position and the main reason for it."}) |
| claims = [p for p in parts if p["type"] == "Claim"] |
| for c in claims: |
| nxt = [p for p in parts if p["start"] > c["start"]] |
| if not nxt or nxt[0]["type"] != "Evidence": |
| out.append({"kind": "part", "start": c["start"], "end": c["end"], |
| "text": "This claim has no evidence right after it: add an example, a fact or a reason."}) |
| for p in parts: |
| if p["effectiveness"] == "Ineffective": |
| out.append({"kind": "part", "start": p["start"], "end": p["end"], |
| "text": f"This {p['type'].lower()} reads as ineffective: make it specific and tie it to your position."}) |
| words = len(text.split()) |
| if words < 250: |
| out.append({"kind": "length", "text": f"{words} words. IELTS Task 2 needs at least 250; short essays lose marks before they are read."}) |
| return out |
|
|
