| import torch |
| from src.models import GrcResult, ProcessGrcInput, TokenConfidence |
|
|
|
|
| def process_grc(model, tokenizer, input: ProcessGrcInput): |
| messages = [ |
| { |
| "content": f"You are a Malay language spelling corrector. I will give you some text written in messy Malay (shortened or mistyped). Rewrite it in correct Malay spelling.\n{input.garbled}", |
| "role": "user", |
| }, |
| {"content": input.gt, "role": "assistant"}, |
| ] |
| top_k = 5 |
|
|
| inputs = tokenizer.apply_chat_template( |
| messages, |
| tokenize=True, |
| add_generation_prompt=False, |
| return_tensors="pt", |
| return_assistant_tokens_mask=True, |
| return_dict=True, |
| ).to("cuda") |
|
|
| attention_mask = inputs.get("attention_mask", None) |
| assistant_mask = inputs.get("assistant_masks", None) |
|
|
| with torch.no_grad(): |
| outputs = model(inputs["input_ids"], attention_mask=attention_mask) |
|
|
| confidences = outputs.logits.softmax(dim=-1) |
|
|
| shifted_confidences = confidences[:, :-1, :] |
| shifted_labels = inputs["input_ids"][:, 1:] |
| shifted_assistant_mask = assistant_mask[:, 1:] |
| boolean_mask = shifted_assistant_mask.bool() |
| labels_on_assistant = shifted_labels[boolean_mask] |
| confs_on_assistant = shifted_confidences[ |
| boolean_mask, : |
| ] |
|
|
| values, indices = torch.topk( |
| confs_on_assistant, k=top_k, dim=-1 |
| ) |
| values = values.to(dtype=torch.float16).cpu().numpy() |
| tokens_decoded: list[list[str]] = [ |
| tokenizer.convert_ids_to_tokens(vals) for vals in indices |
| ] |
|
|
| |
| correct_labels_confidences = ( |
| confs_on_assistant.gather(-1, labels_on_assistant.unsqueeze(1)) |
| .squeeze(1) |
| .to(dtype=torch.float16) |
| .cpu() |
| .numpy() |
| ) |
|
|
| |
| |
| _, sorted_indices = torch.sort(confs_on_assistant, descending=True, dim=-1) |
| correct_labels_ranks = ( |
| torch.nonzero( |
| sorted_indices == labels_on_assistant.unsqueeze(1), as_tuple=False |
| )[:, 1] |
| + 1 |
| ) |
| correct_labels_ranks = correct_labels_ranks.cpu().numpy() |
| decoded_labels = tokenizer.convert_ids_to_tokens(labels_on_assistant) |
|
|
| res = GrcResult( |
| prediction_map=[ |
| TokenConfidence(token=tokens[0], confidence=confidences[0]) |
| for tokens, confidences in zip(tokens_decoded, values) |
| ], |
| gt_map=[ |
| TokenConfidence(token=token, confidence=confidence, rank=rank) |
| for (token, confidence, rank) in zip( |
| decoded_labels, correct_labels_confidences, correct_labels_ranks |
| ) |
| ], |
| topk=[ |
| [ |
| TokenConfidence(token=token, confidence=confidence) |
| for (token, confidence) in zip(rt, rc) |
| ] |
| for (rt, rc) in zip(tokens_decoded, values) |
| ], |
| ) |
| return res |
|
|
