| 216 | return total_val_accuracy / len(val_dataloader) |
| 217 | |
| 218 | def inference(model, tokenizer, prompt, effort_levels): |
| 219 | model.eval() |
| 220 | with torch.no_grad(): |
| 221 | inputs = tokenizer(prompt, return_tensors="pt", max_length=MAX_LENGTH, truncation=True, padding="max_length") |
| 222 | input_ids = inputs["input_ids"].to(device) |
| 223 | attention_mask = inputs["attention_mask"].to(device) |
| 224 | |
| 225 | results = [] |
| 226 | for effort in effort_levels: |
| 227 | effort_tensor = torch.tensor([effort], dtype=torch.float).to(device) |
| 228 | logits = model(input_ids, attention_mask, effort_tensor) |
| 229 | probabilities = F.softmax(logits, dim=1) |
| 230 | predicted_approach_index = torch.argmax(probabilities, dim=1).item() |
| 231 | results.append((APPROACHES[predicted_approach_index], probabilities[0][predicted_approach_index].item())) |
| 232 | |
| 233 | return results |
| 234 | |
| 235 | def main(args): |
| 236 | |