optimise

app.py

@@ -23,29 +23,42 @@ def correct(text):
     model.eval()
 
     text = re.sub(r'\u200d', '<ZWJ>', text)
-    inputs = tokenizer(text, return_tensors='pt', padding='do_not_pad', max_length=1024)
+    inputs = tokenizer(
+        text,
+        return_tensors='pt', 
+        padding='do_not_pad', 
+        max_length=1024
+    )
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
-    with torch.no_grad():
+    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]
+    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)
+
+    return prediction_decoded
+
 
 init()