| |
| |
| |
| import torch, torch.nn as nn, numpy as np, gc, os |
| from kan import KAN |
| from sentence_transformers import SentenceTransformer |
| from dataclasses import dataclass |
| from typing import List |
|
|
| torch.set_num_threads(int(os.environ.get('TORCH_THREADS', '1'))) |
|
|
| class TinyKANDistilled(nn.Module): |
| def __init__(self, input_dim=9, grid=3): |
| super().__init__() |
| self.kan_main = KAN(width=[[input_dim,0],[16,0],[8,0],[1,0]], grid=grid) |
| self.linear_branch = nn.Sequential(nn.Linear(input_dim,8), nn.GELU(), nn.Linear(8,1)) |
| self.gate = nn.Sequential(nn.Linear(input_dim,8), nn.ReLU(), nn.Linear(8,2), nn.Softmax(dim=1)) |
| self.skip = nn.Linear(input_dim, 1) |
| def forward(self, x): |
| k = self.kan_main(x); l = self.linear_branch(x); g = self.gate(x) |
| return (g * torch.cat([k,l],dim=1)).sum(dim=1,keepdim=True) + 0.05*self.skip(x) |
| def get_gate_weights(self, x): |
| with torch.no_grad(): return self.gate(x).cpu().numpy() |
|
|
| @dataclass |
| class EvalResult: |
| score: float; score_pct: float; cefr: str; confidence: float |
| gate_kan: float; gate_linear: float; features: np.ndarray |
| importance: np.ndarray; feedback: str; xai_explanation: str |
|
|
| class KANEvaluator: |
| FEATURE_NAMES = ['cos_MiniLM','cos_MPNet','cos_RoBERTa', |
| '|MiniLM-MPNet|','|MiniLM-RoBERTa|','|MPNet-RoBERTa|', |
| 'mean(cos)','max(cos)','min(cos)'] |
| CEFR_LEVELS = [(0.90,'C2'),(0.80,'C1'),(0.70,'B2'),(0.55,'B1'),(0.40,'A2'),(0.0,'A1')] |
| ENCODERS = {'MiniLM':'sentence-transformers/all-MiniLM-L6-v2', |
| 'MPNet':'sentence-transformers/all-mpnet-base-v2', |
| 'RoBERTa':'sentence-transformers/stsb-roberta-base'} |
|
|
| def __init__(self, model_path, device=None): |
| self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu') |
| self.model = TinyKANDistilled(9,3).to(self.device) |
| self.model.load_state_dict(torch.load(model_path, map_location=self.device, weights_only=False)) |
| self.model.eval() |
| print(f"[Evaluator] TinyKAN loaded ✅ ({sum(p.numel() for p in self.model.parameters()):,} params)") |
| self.encoders = {} |
| for name, path in self.ENCODERS.items(): |
| enc = SentenceTransformer(path, device=self.device) |
| enc.eval() |
| self._quantize(enc, name) |
| self.encoders[name] = enc |
| gc.collect() |
| print("[Evaluator] Encoders loaded ✅ (dynamically quantized, int8)") |
|
|
| def _quantize(self, sentence_transformer, name): |
| """Applies dynamic int8 quantization to the Linear layers of the |
| underlying HF transformer to reduce resident memory on CPU. |
| Falls back silently (keeps full precision) if unsupported.""" |
| try: |
| inner = sentence_transformer[0].auto_model |
| quantized = torch.quantization.quantize_dynamic( |
| inner, {torch.nn.Linear}, dtype=torch.qint8) |
| sentence_transformer[0].auto_model = quantized |
| except Exception as e: |
| print(f"[Evaluator] Quantization skipped for {name}: {e}") |
|
|
| def _features(self, s1, s2): |
| cos = [] |
| for enc in self.encoders.values(): |
| e1 = enc.encode([s1], convert_to_numpy=True) |
| e2 = enc.encode([s2], convert_to_numpy=True) |
| e1 /= np.maximum(np.linalg.norm(e1,axis=1,keepdims=True),1e-8) |
| e2 /= np.maximum(np.linalg.norm(e2,axis=1,keepdims=True),1e-8) |
| cos.append(float((e1*e2).sum())) |
| c1,c2,c3 = cos |
| return np.array([c1,c2,c3,abs(c1-c2),abs(c1-c3),abs(c2-c3),(c1+c2+c3)/3,max(cos),min(cos)],dtype=np.float32) |
|
|
| def _predict(self, feat): |
| x = torch.tensor(feat,dtype=torch.float32).unsqueeze(0).to(self.device) |
| with torch.no_grad(): return float(torch.sigmoid(self.model(x)).cpu().numpy()[0][0]) |
|
|
| def _importance(self, feat): |
| base = self._predict(feat) |
| imp = np.array([abs(base - self._predict(np.where(np.arange(9)==i, 0, feat))) for i in range(9)]) |
| return imp / imp.sum() if imp.sum() > 0 else imp |
|
|
| def _cefr(self, score): |
| for t,l in self.CEFR_LEVELS: |
| if score >= t: return l |
| return 'A1' |
|
|
| STRICT_CATEGORIES = {'grammar', 'conjugation', 'spelling'} |
|
|
| @staticmethod |
| def _normalize(s): |
| import re |
| return re.sub(r"[^\w\s']", '', s.lower()).split() |
|
|
| def _surface_penalty(self, learner, reference, source_sentence=None): |
| """For exercise categories that test a specific grammatical form |
| (not a paraphrase), pure semantic similarity is largely blind to |
| small but critical errors, since a missing function word barely |
| shifts the sentence embedding. Two mechanisms address this: |
| (1) if the answer is essentially identical to the flawed source |
| sentence, no correction was attempted at all — the score is |
| collapsed near zero regardless of semantic closeness; (2) otherwise, |
| a continuous word-overlap ratio against the reference is squared, |
| so partial corrections are penalised more than a linear measure |
| would, while an exact match keeps full credit.""" |
| import difflib |
| lw, rw = self._normalize(learner), self._normalize(reference) |
| if lw == rw: |
| return 1.0 |
| if source_sentence: |
| sw = self._normalize(source_sentence) |
| if lw == sw: |
| return 0.10 |
| ratio = difflib.SequenceMatcher(None, lw, rw).ratio() |
| penalty = ratio ** 2 |
| if len(lw) != len(rw): |
| |
| |
| |
| penalty = min(penalty, 0.5) |
| return max(0.15, penalty) |
|
|
| def evaluate(self, reference, learner, category=None, source_sentence=None): |
| feat = self._features(reference, learner) |
| raw_score = self._predict(feat) |
| penalty = self._surface_penalty(learner, reference, source_sentence) if category in self.STRICT_CATEGORIES else 1.0 |
| score = raw_score * penalty |
| cefr = self._cefr(score) |
| imp = self._importance(feat) |
| x = torch.tensor(feat,dtype=torch.float32).unsqueeze(0).to(self.device) |
| gates = self.model.get_gate_weights(x)[0] |
| top_idx = np.argmax(imp) |
| top_feat = self.FEATURE_NAMES[top_idx] |
|
|
| if score >= 0.90: fb = f"Excellent — Score {score*100:.1f}% | Level {cefr}" |
| elif score >= 0.75: fb = f"Good — Score {score*100:.1f}% | Level {cefr}" |
| elif score >= 0.55: fb = f"Satisfactory — Score {score*100:.1f}% | Level {cefr}" |
| elif score >= 0.40: fb = f"Needs improvement — Score {score*100:.1f}% | Level {cefr}" |
| else: fb = f"Insufficient — Score {score*100:.1f}% | Level {cefr}" |
|
|
| xai = f"Most influential feature: '{top_feat}' ({imp[top_idx]*100:.1f}% importance)" |
| if top_idx < 3: xai += f" — Model relies on {['MiniLM','MPNet','RoBERTa'][top_idx]} cosine similarity." |
| elif top_idx < 6: xai += " — Encoder divergence detected (semantic difficulty signal)." |
| else: xai += f" — {['Mean','Maximum','Minimum'][top_idx-6]} cosine is the decisive statistic." |
| if penalty <= 0.10: |
| xai += f" Note: semantic similarity alone gave {raw_score*100:.1f}%; the answer matches the uncorrected source sentence, so no correction was made — score collapsed accordingly." |
| elif penalty < 1.0: |
| xai += f" Note: semantic similarity alone gave {raw_score*100:.1f}%; partial mismatch against the expected correction reduced the final score." |
|
|
| return EvalResult(score=score, score_pct=score*100, cefr=cefr, |
| confidence=min(min([abs(score-t) for t,_ in self.CEFR_LEVELS])*2,1.0), |
| gate_kan=float(gates[0]), gate_linear=float(gates[1]), |
| features=feat, importance=imp, feedback=fb, xai_explanation=xai) |
