Tokenizes the input document and predicts credential tokens.
(document, model, tokenizer, device="cuda")
| 4 | # model = PeftModel.from_pretrained(base_model, adapter_path) |
| 5 | |
| 6 | def classify_tokens(document, model, tokenizer, device="cuda"): |
| 7 | """Tokenizes the input document and predicts credential tokens.""" |
| 8 | |
| 9 | # Tokenize the document |
| 10 | tokenized = tokenizer(document, padding="max_length", truncation=True, return_offsets_mapping=True, return_tensors="pt") |
| 11 | |
| 12 | # Move input tensors to the correct device |
| 13 | input_ids = tokenized["input_ids"].to(device) |
| 14 | attention_mask = tokenized["attention_mask"].to(device) |
| 15 | offsets = tokenized["offset_mapping"].tolist()[0] # Get token offsets |
| 16 | |
| 17 | # Run inference |
| 18 | with torch.no_grad(): |
| 19 | outputs = model(input_ids, attention_mask=attention_mask) |
| 20 | |
| 21 | # Get predicted labels (0 or 1) |
| 22 | predictions = torch.argmax(outputs.logits, dim=-1).cpu().numpy()[0] # Get labels for first (only) batch |
| 23 | |
| 24 | # Extract credential spans using offsets |
| 25 | credential_tokens = [] |
| 26 | current_credential = "" |
| 27 | previous_end = -1 # Track last token's end position |
| 28 | |
| 29 | for i, label in enumerate(predictions): |
| 30 | if label == 1: # Credential token |
| 31 | start, end = offsets[i] |
| 32 | if start == 0 and end == 0: # Skip special tokens |
| 33 | continue |
| 34 | |
| 35 | token_text = document[start:end] |
| 36 | |
| 37 | if previous_end != -1 and start > previous_end + 1: |
| 38 | # New credential detected, store previous and start a new one |
| 39 | credential_tokens.append(current_credential) |
| 40 | current_credential = token_text |
| 41 | else: |
| 42 | # Append to existing credential |
| 43 | current_credential += token_text |
| 44 | |
| 45 | previous_end = end # Update last token end position |
| 46 | |
| 47 | if current_credential: |
| 48 | credential_tokens.append(current_credential) |
| 49 | |
| 50 | return credential_tokens |