feat: grader info

app.py

@@ -1,13 +1,14 @@
 from typing import *
 import gradio as gr
-from model.IELTS_essay_scoring import grade_IELTS_essay_scoring
-from model.Engessay_grading_ML import grade_Engessay_grading_ML
-from model.mistral_7b_ielts_evaluator import grade_mistral_7b_ielts_evaluator
+from model.common import Grader
+from model.IELTS_essay_scoring import IELTS_essay_scoring
+from model.Engessay_grading_ML import Engessay_grading_ML
+from model.mistral_7b_ielts_evaluator import Mistral_7b_IELTS_Evaluator
 
-models = {
-    "IELTS_essay_scoring": grade_IELTS_essay_scoring,
-    "Engessay_grading_ML": grade_Engessay_grading_ML,
-    "mistral_7b_ielts_evaluator": grade_mistral_7b_ielts_evaluator,
+models: Dict[str, Grader] = {
+    "IELTS_essay_scoring": IELTS_essay_scoring(),
+    "Engessay_grading_ML": Engessay_grading_ML(),
+    "mistral_7b_ielts_evaluator": Mistral_7b_IELTS_Evaluator(),
 }
 
 
@@ -20,7 +21,7 @@ def grade(question: str, answer: str, model: str) -> Tuple[float, str]:
         raise gr.Error(f"Model {model} not found")
 
     grader = models[model]
-    return grader(question, answer)
+    return grader.grade(question, answer)
 
 
 with gr.Blocks() as app:
@@ -46,7 +47,7 @@ with gr.Blocks() as app:
             btn = gr.Button("Grade Essay", variant="primary")
         with gr.Column():
             overall = gr.Number(label="Overall Score")
-            comment = gr.Textbox(label="Comment", lines=10)
+            comment = gr.Markdown(label="Comment")
 
     btn.click(
         fn=grade,
@@ -79,4 +80,7 @@ with gr.Blocks() as app:
         inputs=[question, essay],
     )
 
+    info = "\n\n".join([f"### {k}\n{grader.info()}" for k, grader in models.items()])
+    gr.Markdown(info)
+
 app.launch()

model/Engessay_grading_ML.py

@@ -1,37 +1,42 @@
 from typing import *
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from .common import Grader
 
 model_name = "KevSun/Engessay_grading_ML"
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
-@torch.no_grad()
-def grade_Engessay_grading_ML(question: str, answer: str) -> Tuple[float, str]:
-    text = f"{question} {answer}"
+class Engessay_grading_ML(Grader):
+    def info(self) -> str:
+        return "[KevSun/Engessay_grading_ML](https://huggingface.co/KevSun/Engessay_grading_ML)"
 
-    inputs = tokenizer(text, return_tensors="pt")
+    @torch.no_grad()
+    def grade(self, question: str, answer: str) -> Tuple[float, str]:
+        text = f"{question} {answer}"
 
-    outputs = model(**inputs)
-    predictions = outputs.logits.squeeze()
+        inputs = tokenizer(text, return_tensors="pt")
 
-    predicted_scores = predictions.numpy()
-    scaled_scores = 2.25 * predicted_scores - 1.25
-    rounded_scores = [round(score * 2) / 2 for score in scaled_scores]
+        outputs = model(**inputs)
+        predictions = outputs.logits.squeeze()
 
-    labels = [
-        "cohesion",
-        "syntax",
-        "vocabulary",
-        "phraseology",
-        "grammar",
-        "conventions",
-    ]
-    overall_score = round(sum(rounded_scores) / len(rounded_scores) * 2) / 2
+        predicted_scores = predictions.numpy()
+        scaled_scores = 2.25 * predicted_scores - 1.25
+        rounded_scores = [round(score * 2) / 2 for score in scaled_scores]
 
-    comment = ""
-    for label, score in zip(labels, rounded_scores):
-        comment += f"{label}: {score}\n"
+        labels = [
+            "cohesion",
+            "syntax",
+            "vocabulary",
+            "phraseology",
+            "grammar",
+            "conventions",
+        ]
+        overall_score = round(sum(rounded_scores) / len(rounded_scores) * 2) / 2
 
-    return overall_score, comment
+        comment = ""
+        for label, score in zip(labels, rounded_scores):
+            comment += f"{label}: {score}\n"
+
+        return overall_score, comment

model/IELTS_essay_scoring.py

@@ -2,38 +2,43 @@ from typing import *
 import torch
 import numpy as np
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from .common import Grader
 
 model_name = "JacobLinCool/IELTS_essay_scoring_safetensors"
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
-@torch.no_grad()
-def grade_IELTS_essay_scoring(question: str, answer: str) -> Tuple[float, str]:
-    text = f"{question} {answer}"
+class IELTS_essay_scoring(Grader):
+    def info(self) -> str:
+        return "Safetensors version of [KevSun/IELTS_essay_scoring](https://huggingface.co/KevSun/IELTS_essay_scoring)"
 
-    inputs = tokenizer(
-        text, return_tensors="pt", padding=True, truncation=True, max_length=512
-    )
+    @torch.no_grad()
+    def grade(self, question: str, answer: str) -> Tuple[float, str]:
+        text = f"{question} {answer}"
 
-    outputs = model(**inputs)
-    predictions = outputs.logits.squeeze()
+        inputs = tokenizer(
+            text, return_tensors="pt", padding=True, truncation=True, max_length=512
+        )
 
-    predicted_scores = predictions.numpy()
-    normalized_scores = (predicted_scores / predicted_scores.max()) * 9
-    rounded_scores = np.round(normalized_scores * 2) / 2
+        outputs = model(**inputs)
+        predictions = outputs.logits.squeeze()
 
-    labels = [
-        "Task Achievement",
-        "Coherence and Cohesion",
-        "Vocabulary",
-        "Grammar",
-        "Overall",
-    ]
-    overall_score = float(rounded_scores[-1])
+        predicted_scores = predictions.numpy()
+        normalized_scores = (predicted_scores / predicted_scores.max()) * 9
+        rounded_scores = np.round(normalized_scores * 2) / 2
 
-    comment = ""
-    for label, score in zip(labels, rounded_scores):
-        comment += f"{label}: {score}\n"
+        labels = [
+            "Task Achievement",
+            "Coherence and Cohesion",
+            "Vocabulary",
+            "Grammar",
+            "Overall",
+        ]
+        overall_score = float(rounded_scores[-1])
 
-    return overall_score, comment
+        comment = ""
+        for label, score in zip(labels, rounded_scores):
+            comment += f"{label}: {score}\n"
+
+        return overall_score, comment

model/common.py

@@ -1,4 +1,17 @@
+from typing import Tuple
+import spaces
 from accelerate import Accelerator
 
 accelerator = Accelerator()
 device = accelerator.device
+
+
+class Grader:
+    def info(self) -> str:
+        raise NotImplementedError("Method info not implemented")
+
+    def grade(self, question: str, answer: str) -> Tuple[float, str]:
+        raise NotImplementedError("Method grade not implemented")
+
+    def __call__(self, question: str, answer: str) -> Tuple[float, str]:
+        return self.grade(question, answer)

model/mistral_7b_ielts_evaluator.py

@@ -2,7 +2,7 @@ from typing import *
 import torch
 from transformers import pipeline
 import spaces
-from .common import device
+from .common import Grader, device
 
 model_name = "JacobLinCool/mistral-7b-ielts-evaluator-safetensors"
 pipe = pipeline("text-generation", model_name, device=device)
@@ -44,10 +44,14 @@ system = """In this task, you are required to evaluate an IELTS Writing Task 2 e
 - Suggest strategies for enhancement in weaker areas."""
 
 
-@spaces.GPU(duration=120)
-@torch.no_grad()
-def grade_mistral_7b_ielts_evaluator(question: str, answer: str) -> Tuple[float, str]:
-    text = f"""{system}
+class Mistral_7b_IELTS_Evaluator(Grader):
+    def info(self) -> str:
+        return "Safetensors version of [chillies/mistral-7b-ielts-evaluator-q4](https://huggingface.co/chillies/mistral-7b-ielts-evaluator-q4)"
+
+    @spaces.GPU(duration=120)
+    @torch.no_grad()
+    def grade(self, question: str, answer: str) -> Tuple[float, str]:
+        text = f"""{system}
 
 ## Prompt:
 {question}
@@ -58,10 +62,10 @@ def grade_mistral_7b_ielts_evaluator(question: str, answer: str) -> Tuple[float,
 ## Evaluation:
 """
 
-    outputs = pipe(text, max_length=2048)
-    comment = outputs[0]["generated_text"].split("## Evaluation:")[1].strip()
-    overall_score = float(
-        comment.split("Suggested Overall Band Score: ")[1].split("\n")[0] or 0.0
-    )
+        outputs = pipe(text, max_length=2048)
+        comment = outputs[0]["generated_text"].split("## Evaluation:")[1].strip()
+        overall_score = float(
+            comment.split("Suggested Overall Band Score: ")[1].split("\n")[0] or 0.0
+        )
 
-    return overall_score, comment
+        return overall_score, comment