threads

app.py

@@ -1,65 +1,60 @@
 import gradio as gr
-# import torch
-# import re
+import torch
+import re
 
-# model = None
-# tokenizer = None
+model = None
+tokenizer = None
 
-# def init():
-#     from transformers import MT5ForConditionalGeneration, T5TokenizerFast
-#     import os
-
-#     global model, tokenizer
+def init():
+    from transformers import MT5ForConditionalGeneration, T5TokenizerFast
+    import os
 
-#     hf_token = os.environ.get("HF_TOKEN")
+    global model, tokenizer
 
-#     model = MT5ForConditionalGeneration.from_pretrained("lm-spell/mt5-base-ft-ssc", token=hf_token).to('cpu')
-#     torch.set_num_threads(2)
-#     tokenizer = T5TokenizerFast.from_pretrained("google/mt5-base")
-#     tokenizer.add_special_tokens({'additional_special_tokens': ['<ZWJ>']})
+    hf_token = os.environ.get("HF_TOKEN")
 
+    model = MT5ForConditionalGeneration.from_pretrained("lm-spell/mt5-base-ft-ssc", token=hf_token)
+    torch.set_num_threads(16)
+    tokenizer = T5TokenizerFast.from_pretrained("google/mt5-base")
+    tokenizer.add_special_tokens({'additional_special_tokens': ['<ZWJ>']})
 
-# def correct(text):
 
-#     model.eval()
+def correct(text):
 
-#     text = re.sub(r'\u200d', '<ZWJ>', text)
-#     inputs = tokenizer(
-#         text,
-#         return_tensors='pt',
-#         padding='do_not_pad',
-#         max_length=1024
-#     )
+    model.eval()
 
-#     with torch.inference_mode():
-#         outputs = model.generate(
-#             input_ids=inputs["input_ids"],
-#             attention_mask=inputs["attention_mask"],
-#             max_length=1024,
-#             num_beams=1,
-#             do_sample=False,
-#         )
-#     prediction = outputs[0]
+    text = re.sub(r'\u200d', '<ZWJ>', text)
+    inputs = tokenizer(
+        text,
+        return_tensors='pt',
+        padding='do_not_pad',
+        max_length=1024
+    )
 
-#     special_token_id_to_keep = tokenizer.convert_tokens_to_ids('<ZWJ>')
-#     all_special_ids = set(tokenizer.all_special_ids)
-#     pred_tokens = prediction.cpu()
+    with torch.inference_mode():
+        outputs = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            max_length=1024,
+            num_beams=1,
+            do_sample=False,
+        )
+    prediction = outputs[0]
 
-#     tokens_list = pred_tokens.tolist()
-#     filtered_tokens = [
-#         token for token in tokens_list 
-#         if token == special_token_id_to_keep or token not in all_special_ids
-#     ]
+    special_token_id_to_keep = tokenizer.convert_tokens_to_ids('<ZWJ>')
+    all_special_ids = set(tokenizer.all_special_ids)
+    pred_tokens = prediction.cpu()
 
-#     prediction_decoded = tokenizer.decode(filtered_tokens, skip_special_tokens=False).replace('\n', '').strip()
+    tokens_list = pred_tokens.tolist()
+    filtered_tokens = [
+        token for token in tokens_list 
+        if token == special_token_id_to_keep or token not in all_special_ids
+    ]
 
-#     return re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)
+    prediction_decoded = tokenizer.decode(filtered_tokens, skip_special_tokens=False).replace('\n', '').strip()
 
-def test(text):
-    import os
-    return f"vCPUs: {os.cpu_count()}"
+    return re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)
 
-# init()
-# demo = gr.Interface(fn=correct, inputs="text", outputs="text")
-demo = gr.Interface(fn=test, inputs="text", outputs="text")
+init()
+demo = gr.Interface(fn=correct, inputs="text", outputs="text")
 demo.launch()