Update app.py

app.py

@@ -8,9 +8,9 @@ from   json import loads
 checkpoint = "IsmaelMousa/SmolLM2-135M-Instruct-EngSaf-217K"
 device     = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model     = AutoModelForCausalLM.from_pretrained(checkpoint)
-assistant = pipeline("text-generation", tokenizer=tokenizer, model=model, device=device)
+tokenizer  = AutoTokenizer.from_pretrained(checkpoint)
+model      = AutoModelForCausalLM.from_pretrained(checkpoint)
+assistant  = pipeline("text-generation", tokenizer=tokenizer, model=model, device=device)
 
 
 def extract(text):
@@ -26,7 +26,7 @@ def extract(text):
     except: return response
 
 
-def grade(question, reference_answer, student_answer, mark_scheme, max_new_tokens):
+def grade(question, reference_answer, student_answer, mark_scheme):
     system_content = "You are a grading assistant. Evaluate student answers based on the mark scheme. Respond only in JSON format with keys \"score\" (int) and \"rationale\" (string)."
 
     user_content   = ("Provide both a score and a rationale by evaluating the student's answer strictly within the mark scheme range, "
@@ -40,7 +40,7 @@ def grade(question, reference_answer, student_answer, mark_scheme, max_new_token
     
     inputs         = tokenizer.apply_chat_template(messages, tokenize=False)
 
-    output         = assistant(inputs, max_new_tokens=max_new_tokens, do_sample=False, return_full_text=False)[0]["generated_text"]
+    output         = assistant(inputs, max_new_tokens=128, do_sample=False, return_full_text=False)[0]["generated_text"]
     parsed         = extract(output)
     
     return parsed
@@ -51,8 +51,7 @@ demo = gr.Interface(fn    =grade,
                             gr.Textbox(label="Reference Answer"),
                             gr.Textbox(label="Student Answer"),
                             gr.Textbox(label="Mark Scheme"),
-                            gr.Slider(minimum=1, maximum=512, value=128, step=1, label="Max new tokens")],
-                    outputs=gr.JSON(label="Evaluation Output"))
+                    outputs=gr.JSON   (label="Evaluation Output"))
 
 
 if __name__ == "__main__": demo.launch()