Update README.md

README.md

@@ -98,41 +98,33 @@ from transformers import GenerationConfig, AutoTokenizer
 import torch
 
 model_id = "TheBloke/Mistral-7B-Instruct-v0.2-GPTQ"
-adapter="nmarafo/Mistral-7B-Instruct-v0.2-TrueFalse-Feedback-GPTQ"
+adapter = "nmarafo/Mistral-7B-Instruct-v0.2-TrueFalse-Feedback-GPTQ"
 
-# To perform inference on the test dataset example load the model from the checkpoint
-persisted_model = AutoPeftModelForCausalLM.from_pretrained(
-    adapter,
-    low_cpu_mem_usage=True,
-    return_dict=True,
-    torch_dtype=torch.float16,
-    device_map="cuda")
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, return_token_type_ids=False)
+tokenizer.pad_token = tokenizer.eos_token
 
-def generate_prompt(data_point):
-    system_message = "Analyze the question, the expected answer, and the student's response. Determine if the student's answer is conceptually correct in relation to the expected answer, regardless of the exact wording. An answer will be considered correct if it accurately identifies the key information requested in the question, even if expressed differently. Return True if the student's answer is correct or False otherwise. Add a brief comment explaining the rationale behind the answer being correct or incorrect."
-    question = data_point["question"][0]
-    best_answer = data_point["best_answer"][0]
-    student_answer = data_point["student_answer"][0]
-    prompt = f"{system_message}\n\nQuestion: {question}\nExpected Answer: {best_answer}\nStudent Answer: {student_answer}"
+model = AutoPeftModelForCausalLM.from_pretrained(adapter, low_cpu_mem_usage=True, return_dict=True, torch_dtype=torch.float16, device_map="cuda")
 
-    return prompt
+def predict(question, best_answer, student_answer):
+    system_message = "Analyze the question, the expected answer, and the student's response. Determine if the student's answer is conceptually correct in relation to the expected answer, regardless of the exact wording. Return True if the student's answer is correct or False otherwise. Add a brief comment explaining the rationale behind the answer being correct or incorrect."
+    prompt = f"{system_message}\n\nQuestion: {question}\nBest Answer: {best_answer}\nStudent Answer: {student_answer}"
+    prompt_template=f"<s>[INST]{prompt}[/INST]"
 
-tokenizer = AutoTokenizer.from_pretrained(
-            model_id, 
-            trust_remote_code=True, 
-            return_token_type_ids=False)
-tokenizer.pad_token = tokenizer.eos_token
+    encoding = tokenizer(prompt_template, return_tensors='pt', padding=True, truncation=True, max_length=512)
+    input_ids = encoding['input_ids'].cuda()
+    attention_mask = encoding['attention_mask'].cuda()
 
-question="Name of Canary Island"
+    output = model.generate(input_ids, attention_mask=attention_mask, 
+                            temperature=0.7, do_sample=True, top_p=0.95, 
+                            top_k=40, max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
+
+question="Mention all the Canary Island"
 best_answer="Tenerife, Fuerteventura, Gran Canaria, Lanzarote, La Palma, La Gomera, El Hierro, La Graciosa"
 student_answer="Tenerife"
 
-prompt = generate_prompt([{"question":question, "best_answer":best_answer,"student_answer":student_answer}])
-prompt_template=f'''<s>[INST] {prompt} [/INST]'''
-
-input_ids = tokenizer(prompt, return_tensors='pt').input_ids.cuda()
-output = persisted_model.generate(inputs=input_ids, temperature=0.7, do_sample=True, top_p=0.95, top_k=40, max_new_tokens=512)
-print(tokenizer.decode(output[0]))    
+print(predict(question, best_answer, student_answer))    
 
 ```