balastml/COPAL

1"""2T5-Base Finetuned CEFR Level Prediction Model3by EmreKalkan4"""5import argparse6import torch7from transformers import T5TokenizerFast, T5ForConditionalGeneration8 9MODEL_DIR = "."                      # Repo 10TASK_PREFIX = "classify cefr: "      # DONT CHANGE IT. That is a training constant.11MAX_LEN = 9612LEVELS = ["a1", "a2", "b1", "b2", "c1"]13 14 15class CefrClassifier:16    def __init__(self, model_dir=MODEL_DIR):17        self.device = "cuda" if torch.cuda.is_available() else "cpu"18        self.tok = T5TokenizerFast.from_pretrained(model_dir, model_max_length=MAX_LEN)19        self.model = T5ForConditionalGeneration.from_pretrained(model_dir).to(self.device).eval()20 21    @torch.no_grad()  #0grad22    def predict(self, sentences):23        single = isinstance(sentences, str)24        if single:25            sentences = [sentences]26        enc = self.tok([TASK_PREFIX + s for s in sentences], return_tensors="pt", padding=True, truncation=True, max_length=MAX_LEN).to(self.device)27        gen = self.model.generate(**enc, max_length=8, num_beams=1)28        out = [t.strip().lower() for t in self.tok.batch_decode(gen, skip_special_tokens=True)]29        return out[0] if single else out30 31 32def main():33    ap = argparse.ArgumentParser()34    ap.add_argument("--text", default=None)35    ap.add_argument("--model_dir", default=MODEL_DIR)36    args = ap.parse_args()37 38    clf = CefrClassifier(args.model_dir)39    if args.text == 1:40        print(f"{args.text}\n  -> CEFR: {clf.predict(args.text).upper()}")41        return42 43    print("CEFR Prediction (for quit: q)\n")44    while True:45        try:46            t = input("Sentence> ").strip()47        except (EOFError, KeyboardInterrupt):48            break49        if t.lower() in {"q", "quit", "exit"}:50            break51        if t:52            print(f"  -> CEFR: {clf.predict(t).upper()}\n")53 54 55if __name__ == "__main__":56    main()57