zwj fix

app.py

@@ -25,8 +25,8 @@ def correct(text):
     text = re.sub(r'\u200d', '<ZWJ>', text)
     inputs = tokenizer(
         text,
-        return_tensors='pt', 
-        padding='do_not_pad', 
+        return_tensors='pt',
+        padding='do_not_pad',
         max_length=1024
     )
     inputs = {k: v.to(device) for k, v in inputs.items()}
@@ -41,23 +41,23 @@ def correct(text):
         )
     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)
+    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()
+    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()
+    prediction_decoded = tokenizer.decode(filtered_tokens, skip_special_tokens=False).replace('\n', '').strip()
     
-    # return re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)
+    return re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)
 
-    prediction_decoded = tokenizer.decode(prediction, skip_special_tokens=True).replace('\n', '').strip()
-    prediction_decoded = re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)
+    # prediction_decoded = tokenizer.decode(prediction, skip_special_tokens=True).replace('\n', '').strip()
+    # prediction_decoded = re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)
 
-    return prediction_decoded
+    # return prediction_decoded
 
 
 init()