wip

app.py

@@ -42,18 +42,32 @@ 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)
+    all_special_ids = set(tokenizer.all_special_ids)
+    pred_tokens = prediction.cpu()
+
+    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
+    ]
 
-    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)
 
+    # 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)
+
     # prediction_decoded = tokenizer.decode(prediction, skip_special_tokens=True).replace('\n', '').strip()
     # prediction_decoded = re.sub(r'<ZWJ>\s?', '\u200d', prediction_decoded)