Update app.py

app.py

@@ -4,7 +4,7 @@ import sys
 import dataset
 import engine
 from model import BERTBaseUncased
-# from tokenizer import tokenizer
+from tokenizer import tokenizer
 import config
 from transformers import pipeline, AutoTokenizer, AutoModel
 import gradio as gr
@@ -14,32 +14,32 @@ model = BERTBaseUncased()
 model.load_state_dict(torch.load(config.MODEL_PATH, map_location=torch.device(device)),strict=False)
 model.to(device)
 
-# T = tokenizer.TweetTokenizer(
-#     preserve_handles=True, preserve_hashes=True, preserve_case=False, preserve_url=False)
-
-# def preprocess(text):
-#     tokens = T.tokenize(text)
-#     print(tokens, file=sys.stderr)
-#     ptokens = []
-#     for index, token in enumerate(tokens):
-#         if "@" in token:
-#             if index > 0:
-#                 # check if previous token was mention
-#                 if "@" in tokens[index-1]:
-#                     pass
-#                 else:
-#                     ptokens.append("mention_0")
-#             else:
-#                 ptokens.append("mention_0")
-#         else:
-#             ptokens.append(token)
-
-#     print(ptokens, file=sys.stderr)
-#     return " ".join(ptokens)
+T = tokenizer.TweetTokenizer(
+    preserve_handles=True, preserve_hashes=True, preserve_case=False, preserve_url=False)
+
+def preprocess(text):
+    tokens = T.tokenize(text)
+    print(tokens, file=sys.stderr)
+    ptokens = []
+    for index, token in enumerate(tokens):
+        if "@" in token:
+            if index > 0:
+                # check if previous token was mention
+                if "@" in tokens[index-1]:
+                    pass
+                else:
+                    ptokens.append("mention_0")
+            else:
+                ptokens.append("mention_0")
+        else:
+            ptokens.append(token)
+
+    print(ptokens, file=sys.stderr)
+    return " ".join(ptokens)
 
 
 def sentence_prediction(sentence):
-    # sentence = preprocess(sentence)
+    sentence = preprocess(sentence)
     
     model_path = config.MODEL_PATH
 
@@ -51,7 +51,7 @@ def sentence_prediction(sentence):
     test_data_loader = torch.utils.data.DataLoader(
         test_dataset,
         batch_size=config.VALID_BATCH_SIZE,
-        num_workers=-1
+        num_workers=2
     )
 
     outputs, [] = engine.predict_fn(test_data_loader, model, device)