File size: 3,700 Bytes
e1679fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Fit the IELTS band estimate on top of the two trained models, on the ~1.2k human-banded IELTS essays.
MODAL_PROFILE=labs3 modal run scripts/band.py --scores scores-v1 --parts parts-v1
Features per essay: the six raw scores, word count, number of parts of each type, share of effective parts. Model:
ridge regression, 5-fold cross-validation, reported as exact-band and within-half-band accuracy and RMSE in bands.
The fitted weights go to /data/runs/band-v1 (band.json) so inference is a dot product.
"""
import modal
app = modal.App("essay-coach-band")
vol = modal.Volume.from_name("essay-data")
cache = modal.Volume.from_name("essay-hf-cache", create_if_missing=True)
image = (modal.Image.debian_slim(python_version="3.11")
.pip_install("torch==2.4.1", index_url="https://download.pytorch.org/whl/cu124")
.pip_install("transformers==4.46.3", "sentencepiece", "protobuf", "tiktoken", "safetensors", "pandas", "numpy<2.3",
"scikit-learn", "huggingface_hub")
.add_local_python_source("coach"))
@app.function(image=image, gpu="A10G", volumes={"/data": vol, "/root/.cache/huggingface": cache}, timeout=3600)
def fit(scores: str, parts: str, tag: str):
import json
import os
import numpy as np
import pandas as pd
from sklearn.linear_model import Ridge
from sklearn.model_selection import KFold
from coach.data import TRAITS, TYPES
from coach.infer import Coach
df = pd.read_csv("/data/ielts-writing-scored-essays-dataset/ielts_writing_dataset.csv")
df = df[df.Overall.notna() & df.Essay.notna()].copy()
df["Task_Type"] = df.Task_Type.astype(int)
print(f"IELTS essays with a band: {len(df)} (task 2: {(df.Task_Type == 2).sum()})", flush=True)
coach = Coach(f"/data/runs/{scores}", f"/data/runs/{parts}")
feats = []
for i, r in enumerate(df.itertuples(index=False)):
read = coach.read(str(r.Essay))
raw = read["scores_raw"]
ps = read["parts"]
f = [raw[t] for t in TRAITS] + [np.log1p(len(str(r.Essay).split()))] + [sum(p["type"] == t for p in ps) for t in TYPES]
f += [np.mean([p["effectiveness"] == "Effective" for p in ps]) if ps else 0.0, float(r.Task_Type == 2)]
feats.append(f)
if i % 200 == 0:
print(i, flush=True)
X, y = np.array(feats, dtype=np.float32), df.Overall.values.astype(np.float32)
names = TRAITS + ["log_words"] + [f"n_{t}" for t in TYPES] + ["share_effective", "task2"]
mu, sd = X.mean(0), X.std(0) + 1e-6
Xn = (X - mu) / sd
pred = np.zeros_like(y)
for trn, val in KFold(5, shuffle=True, random_state=0).split(Xn):
m = Ridge(alpha=3.0).fit(Xn[trn], y[trn])
pred[val] = m.predict(Xn[val])
snapped = np.clip(np.round(pred * 2) / 2, 1, 9)
res = {"n": int(len(y)), "rmse_bands": float(np.sqrt(((pred - y) ** 2).mean())),
"exact": float((snapped == y).mean()), "within_half": float((np.abs(snapped - y) <= 0.5).mean()),
"within_one": float((np.abs(snapped - y) <= 1.0).mean()), "baseline_rmse_mean": float(y.std()),
"corr": float(np.corrcoef(pred, y)[0, 1])}
print(res, flush=True)
m = Ridge(alpha=3.0).fit(Xn, y)
run = f"/data/runs/{tag}"
os.makedirs(run, exist_ok=True)
json.dump({"features": names, "mean": mu.tolist(), "std": sd.tolist(), "coef": m.coef_.tolist(),
"intercept": float(m.intercept_), "cv": res, "scores_run": scores, "parts_run": parts}, open(f"{run}/band.json", "w"), indent=1)
vol.commit()
return res
@app.local_entrypoint()
def main(scores: str = "scores-v1", parts: str = "parts-v1", tag: str = "band-v1"):
print(fit.remote(scores, parts, tag))
|