wip

app.py

@@ -1,16 +1,46 @@
 import gradio as gr
+import torch
+import re
 
-x = ''
+model = None
+tokenizer = None
+device = "cpu"
 
 def init():
-    global x
-    x = 'Karu'
-    print("Model or environment initialized.")
+    from transformers import MT5ForConditionalGeneration, T5TokenizerFast
+    global model, tokenizer
 
-def correct(name):
-    return "Hello " + name  + x + "!!"
+    model_path = "lm-spell/mt5-base-ft-ssc"
+    model = MT5ForConditionalGeneration.from_pretrained(model_path).to(device)
+    tokenizer = T5TokenizerFast.from_pretrained("google/mt5-base")
+    tokenizer.add_special_tokens({'additional_special_tokens': ['<ZWJ>']})
+        
+
+def correct(text):
+    model.eval()
+
+    text = re.sub(r'\u200d', '<ZWJ>', text)
+    inputs = tokenizer(text, return_tensors='pt', padding='max_length', truncation=True, max_length=128)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
+        prediction = outputs[0]
+
+    special_token_id_to_keep = tokenizer.convert_tokens_to_ids('<ZWJ>')
+    all_special_ids = torch.tensor(tokenizer.all_special_ids, dtype=torch.int64).to(device)
+    special_token_tensor = torch.tensor([special_token_id_to_keep], dtype=torch.int64).to(device)
+
+    pred_tokens = prediction.to(device)
+    tokens_tensor = pred_tokens.clone().detach().to(dtype=torch.int64)
+    mask = (tokens_tensor == special_token_tensor) | (~torch.isin(tokens_tensor, all_special_ids))
+    filtered_tokens = tokens_tensor[mask].tolist()
+    
+    prediction_decoded = tokenizer.decode(filtered_tokens, skip_special_tokens=False).replace('\n', '').strip()
+    
+    return re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)
 
 init()
 
-demo = gr.Interface(fn=correct, inputs="text", outputs="text")
+demo = gr.Interface(fn=correct, inputs="text", outputs="text", share=True)
 demo.launch()

requirements.txt

@@ -1 +1,3 @@
-gradio
+gradio
+torch==2.5.1
+transformers==4.51.3