Views
No views yet
1import torch
2from transformers import AutoModelForMaskedLM, AutoTokenizer
3
4MODEL_NAME = 'udmurtNLP/bert-tiny-char-ctc-udm-denoise'
5model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME)
6tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
7
8def fix_text(text, verbose=False, spaces=2):
9 with torch.inference_mode():
10 batch = tokenizer(text, return_tensors='pt', spaces=spaces, padding=True, truncation=True, return_token_type_ids=False).to(model.device)
11 logits = torch.log_softmax(model(**batch).logits, axis=-1)
12 decoded = tokenizer.decode(logits[0].argmax(-1), skip_special_tokens=True)
13 return tokenizer.clean_up_tokenization(decoded)
14fix_text("кыче мои солы оскылй!")
15# Кыӵе мон солы оскылӥ!