Update app.py

app.py

@@ -1,15 +1,72 @@
 import gradio as gr
 from transformers import pipeline
 
+from model import BERTBaseUncased
+from tokenizer import tokenizer
+import torch
+from utils import label_full_decoder
+import sys
+import config
+import dataset
+import engine
 from model import BERTBaseUncased
 
+MODEL = None
+DEVICE = config.device
 
 def get_sentiment(input_text):
     result = sentiment(input_text)
     return f"result: {result[0]['label']}", f"score: {result[0]['score']}"
 
+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)
+    model_path = config.MODEL_PATH
+
+    test_dataset = dataset.BERTDataset(
+        review=[sentence],
+        target=[0]
+    )
+
+    test_data_loader = torch.utils.data.DataLoader(
+        test_dataset,
+        batch_size=config.VALID_BATCH_SIZE,
+        num_workers=3
+    )
+
+    device = config.device
+
+    model = BERTBaseUncased()
+    model.load_state_dict(torch.load(
+        model_path, map_location=torch.device(device)))
+    model.to(device)
+
+    outputs, [] = engine.predict_fn(test_data_loader, model, device)
+    print(outputs)
+    return outputs[0]
+    
 interface = gr.Interface(
-    fn=get_sentiment,
+    fn=sentence_prediction,
     inputs='text',
     outputs=['text', 'text'],
     title='Sentiment Analysis',