Update app.py

app.py

@@ -1,15 +1,15 @@
 import numpy as np
 import gradio as gr
+from imports import *
 from huggingface_hub import login
 login(token="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")
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dict_ = {
     0: "negative",
     1: "positive",
     2: "neutral"}
-path_sent = "/content/drive/MyDrive/company-review-analysis-model/sentiment_checkpoint"
-tokenizer_sent = AutoTokenizer.from_pretrained(path_sent, use_fast=False)
-model_sent = AutoModelForSequenceClassification.from_pretrained(path_sent, num_labels=3).to(device)
+tokenizer_sent = AutoTokenizer.from_pretrained("nam194/sentiment", use_fast=False)
+model_sent = AutoModelForSequenceClassification.from_pretrained("nam194/sentiment", num_labels=3, use_auth_token=True).to(device)
 def cvt2cls(data):
     data = list(set(data))
     try: 
@@ -24,20 +24,18 @@ def cvt2cls(data):
     return data
 ner_tags = {0: 'B-chỗ để xe', 1: 'B-con người', 2: 'B-công việc', 3: 'B-cơ sở vật chất', 4: 'B-dự án', 5: 'B-lương', 6: 'B-môi trường làm việc', 7: 'B-ot/thời gian', 8: 'B-văn phòng', 9: 'B-đãi ngộ', 10: 'I-chỗ để xe', 11: 'I-con người', 12: 'I-công việc', 13: 'I-cơ sở vật chất', 14: 'I-dự án', 15: 'I-lương', 16: 'I-môi trường làm việc', 17: 'I-ot/thời gian', 18: 'I-văn phòng', 19: 'I-đãi ngộ', 20: 'O'}
 topic_tags = {0: 'chỗ để xe', 1: 'con người', 2: 'công việc', 3: 'cơ sở vật chất', 4: 'dự án', 5: 'lương', 6: 'môi trường làm việc', 7: 'ot/thời gian', 8: 'văn phòng', 9: 'đãi ngộ'}
-path_topic = "/content/drive/MyDrive/company-review-analysis-model/topic_checkpoint"
-num_labels = 20
-config = RobertaConfig.from_pretrained(path_topic, num_labels=num_labels)
-tokenizer_topic = AutoTokenizer.from_pretrained(path_topic, use_fast=False) 
-model_topic = PhoBertLstmCrf.from_pretrained(path_topic, config=config, from_tf=False).to(device)
+config = RobertaConfig.from_pretrained("nam194/ner", num_labels=20)
+tokenizer_topic = AutoTokenizer.from_pretrained("nam194/ner", use_fast=False) 
+model_topic = PhoBertLstmCrf.from_pretrained("nam194/ner", config=config, from_tf=False).to(device)
 model_topic.resize_token_embeddings(len(tokenizer_topic)) 
 
 
-def review_company(sent: str):
+def sentiment(sent: str):
     try:
-        sent = normalize(text=sent) # segment input sentence, maybe raise ConnectionError: HTTPConnectionPool())
+        sent_ = normalize(text=sent_) # segment input sentence, maybe raise ConnectionError: HTTPConnectionPool())
     except:
         pass
-    input_sent = torch.tensor([tokenizer_sent.encode(sent)]).to(device)
+    input_sent = torch.tensor([tokenizer_sent.encode(sent_)]).to(device)
     with torch.no_grad():
         out_sent = model_sent(input_sent)
         logits_sent = out_sent.logits.softmax(dim=-1).tolist()[0]
@@ -55,12 +53,9 @@ def review_company(sent: str):
         for idx, batch in enumerate(dump_iter):
             batch = { k:v.to(device) for k, v in batch.items() }        
             outputs = model_topic(**batch)
-    return list(set([topic_tags[i] for i in cvt2cls(outputs["tags"][0])]))
-
-    return pred_sent
-
-
-
+    pred_topic = list(set([topic_tags[i] for i in cvt2cls(outputs["tags"][0])]))
+    return {"sentiment": pred_sent, "topic": pred_topic}
+    
 
 def flip_image(x):
     return np.fliplr(x)
@@ -71,17 +66,17 @@ with gr.Blocks() as demo:
     with gr.Tab("Review Company"):
         text_input = gr.Textbox()
         text_output = gr.Textbox()
-        text_button = gr.Button("Flip")
-    with gr.Tab("Flip Image"):
+        text_button = gr.Button("Predict")
+    with gr.Tab("Extract infomation from resume"):
         with gr.Row():
             image_input = gr.Image()
             image_output = gr.Image()
-        image_button = gr.Button("Flip")
+        image_button = gr.Button("Predict")
 
-    with gr.Accordion("Open for More!"):
-        gr.Markdown("Look at me...")
+    # with gr.Accordion("Open for More!"):
+    #     gr.Markdown("Look at me...")
 
-    text_button.click(flip_text, inputs=text_input, outputs=text_output)
+    text_button.click(sentiment, inputs=text_input, outputs=text_output)
     image_button.click(flip_image, inputs=image_input, outputs=image_output)
 
 demo.launch()