Views
No views yet
1from peft import PeftModel, PeftConfig
2from transformers import AutoModelForTokenClassification
3
4config = PeftConfig.from_pretrained("bite-the-byte/byt5-small-deASCIIfy-TR")
5model = AutoModelForTokenClassification.from_pretrained("google/byt5-small")
6model = PeftModel.from_pretrained(model, "bite-the-byte/byt5-small-deASCIIfy-TR")
7
8def test_mask(device, sample):
9 """
10 Masks the padded tokens in the input.
11 Args:
12 data (list): List of strings.
13 Returns:
14 dataset (list): List of dictionaries.
15 """
16
17 tokens = dict()
18
19 input_tokens = [i + 3 for i in sample.encode('utf-8')]
20 input_tokens.append(0) # eos token
21 tokens['input_ids'] = torch.tensor([input_tokens], dtype=torch.int64, device=device)
22
23 # Create attention mask
24 tokens['attention_mask'] = torch.ones_like(tokens['input_ids'], dtype=torch.int64, device=device)
25
26 return tokens
27
28def rewrite(model, data):
29 """
30 Rewrites the input text with the model.
31 Args:
32 model (torch.nn.Module): Model.
33 data (dict): Dictionary containing 'input_ids' and 'attention_mask'.
34 Returns:
35 output (str): Rewritten text.
36 """
37
38 with torch.no_grad():
39 pred = torch.argmax(model(**data).logits, dim=2).squeeze(0)
40
41 output = list() # save the indices of the characters as list of integers
42
43 # Conversion table for Turkish characters {100: [300, 350], ...}
44 en2tr = {en: tr for tr, en in zip(list(map(list, map(str.encode, list('ÜİĞŞÇÖüığşçö')))), list(map(ord, list('UIGSCOuigsco'))))}
45
46 for inp, lab in zip((data['input_ids'].squeeze(0) - 3).tolist(), pred.tolist()):
47 if lab and inp in en2tr:
48 # if the model predicts a diacritic, replace it with the corresponding Turkish character
49 output.extend(en2tr[inp])
50 elif inp >= 0: output.append(inp)
51 return bytes(output).decode()
52
53def try_it(text, model):
54 sample = test_mask(model.device, text)
55 return rewrite(model, sample)
56
57try_it('Cekoslovakyalilastiramadiklarimizdan misiniz?', model)