Create README.md

README.md

@@ -0,0 +1,41 @@
+```python
+from transformers import EncoderDecoderModel
+from importlib.machinery import SourceFileLoader
+from transformers.file_utils import cached_path, hf_bucket_url
+import torch
+import os
+
+## Load model & tokenizer
+cache_dir='./cache'
+model_name='nguyenvulebinh/spelling-oov'
+
+def download_tokenizer_files():
+    resources = ['envibert_tokenizer.py', 'dict.txt', 'sentencepiece.bpe.model']
+    for item in resources:
+        if not os.path.exists(os.path.join(cache_dir, item)):
+            tmp_file = hf_bucket_url(model_name, filename=item)
+            tmp_file = cached_path(tmp_file,cache_dir=cache_dir)
+            os.rename(tmp_file, os.path.join(cache_dir, item))
+        
+download_tokenizer_files()
+spell_tokenizer = SourceFileLoader("envibert.tokenizer",os.path.join(cache_dir,'envibert_tokenizer.py')).load_module().RobertaTokenizer(cache_dir)
+spell_model = EncoderDecoderModel.from_pretrained(model_name)
+
+def oov_spelling(word, num_candidate=1):
+    result = []
+    inputs = spell_tokenizer([word])
+    input_ids = inputs['input_ids']
+    attention_mask = inputs['attention_mask']
+    inputs = {
+        "input_ids": torch.tensor(input_ids),
+        "attention_mask": torch.tensor(attention_mask)
+    }
+    outputs = spell_model.generate(**inputs, num_return_sequences=num_candidate)
+    for output in outputs.cpu().detach().numpy().tolist():
+        result.append(spell_tokenizer.sp_model.DecodePieces(spell_tokenizer.decode(output, skip_special_tokens=True).split()))
+    return result    
+    
+oov_spelling('spacespeaker')
+# output: ['x pây x pếch cơ']
+
+```