| 76 | return encoding['input_ids'], encoding['attention_mask'] |
| 77 | |
| 78 | def predict_approach(model, input_ids, attention_mask, device, effort=0.7): |
| 79 | model.eval() |
| 80 | with torch.no_grad(): |
| 81 | input_ids = input_ids.to(device) |
| 82 | attention_mask = attention_mask.to(device) |
| 83 | effort_tensor = torch.tensor([effort], dtype=torch.float).to(device) |
| 84 | |
| 85 | logits = model(input_ids, attention_mask=attention_mask, effort=effort_tensor) |
| 86 | probabilities = F.softmax(logits, dim=1) |
| 87 | predicted_approach_index = torch.argmax(probabilities, dim=1).item() |
| 88 | confidence = probabilities[0][predicted_approach_index].item() |
| 89 | |
| 90 | return APPROACHES[predicted_approach_index], confidence |
| 91 | |
| 92 | def run(system_prompt, initial_query, client, model, **kwargs): |
| 93 | try: |