Views
No views yet
1import torch
2from transformers import CanineTokenizer, CanineForTokenClassification
3
4tokenizer = CanineTokenizer.from_pretrained('slone/canine-c-bashkir-gec-v1')
5model = CanineForTokenClassification.from_pretrained('slone/canine-c-bashkir-gec-v1')
6if torch.cuda.is_available():
7 model.cuda()
8
9LABELS_THIS = [c[5:] for c in model.config.id2label.values() if c.startswith('THIS_')]
10LABELS_NEXT = [c[5:] for c in model.config.id2label.values() if c.startswith('NEXT_')]
11
12def fix_text(text, boost=0):
13 """Apply the model to edit the text. `boost` is a parameter to control edit aggressiveness."""
14 bx = tokenizer(text, return_tensors='pt', padding=True)
15 with torch.inference_mode():
16 out = model(**bx.to(model.device))
17 n1, n2 = len(LABELS_THIS), len(LABELS_NEXT)
18 logits1 = out.logits[0, :, :n1].view(-1, n1)
19 logits2 = out.logits[0, :, n1:].view(-1, n2)
20 if boost:
21 logits1[1:, 0] -= boost
22 logits2[:, 0] -= boost
23 ids1, ids2 = logits1.argmax(-1).tolist(), logits2.argmax(-1).tolist()
24 result = []
25 for c, id1, id2 in zip(' ' + text, ids1, ids2):
26 l1, l2 = LABELS_THIS[id1], LABELS_NEXT[id2]
27 if l1 == 'KEEP':
28 result.append(c)
29 elif l1 != 'DELETE':
30 result.append(l1)
31 if l2 != 'PASS':
32 result.append(l2)
33 return ''.join(result)
34
35text = 'У йыл дан д ың йөҙө һoрөмлэнде.'
36print(fix_text(text)) # Уйылдандың йөҙө һөрөмләнде.boost can be used to control the aggressiveness of editing:
positive values increase the probability of changing the text, negative values decrease it.