Update app.py

app.py

@@ -1,86 +1,76 @@
 import gradio as gr
 from transformers import pipeline
+import torch
+from openprompt.plms import load_plm
 from openprompt import PromptDataLoader
 from openprompt.prompts import ManualVerbalizer
 from openprompt.prompts import ManualTemplate
+from openprompt.data_utils import InputExample
 from openprompt import PromptForClassification
 
-unmasker = pipeline('fill-mask', model="CCCC/ARCH_tuned_bert") #'bert-base-uncased')
 
-def fill_mask(text):
-  unmasked = unmasker(text)
-  output = {}
-  for unmask in unmasked:
-    output[unmask["token_str"]] = unmask["score"]
-  return output
+def sentiment_analysis(sentence, template, positive, neutral, negative):
+    model_name = "CCCC/ARCH_tuned_bert" 
+    template = template.replace('[SENTENCE]', '{"placeholder":"text_a"}')
+    template = template.replace('[MASK]', '{"mask"}')
+    classes = ['positive', 'neutral', 'negative']
 
-examples = [["Hello I'm a [MASK] model."], ["[MASK] is my favourite sports."]]
+    label_words = {
+        "positive": positive.split(" "),
+        "neutral": neutral.split(" "),
+        "negative": negative.split(" "),
+    }
+    print(label_words)
+    type_dic = {
+        "CCCC/ARCH_tuned_bert":"bert",
+        "bert-base-uncased":"bert",
+        "roberta-base":"roberta",
+        "yiyanghkust/finbert-pretrain":"bert",
+        "facebook/opt-125m":"opt",
+        "facebook/opt-350m":"opt",
+    }
 
-css = """
-footer {display:none !important}
-.output-markdown{display:none !important}
-.gr-button-primary {
-    z-index: 14;
-    height: 43px;
-    width: 130px;
-    left: 0px;
-    top: 0px;
-    padding: 0px;
-    cursor: pointer !important; 
-    background: none rgb(17, 20, 45) !important;
-    border: none !important;
-    text-align: center !important;
-    font-family: Poppins !important;
-    font-size: 14px !important;
-    font-weight: 500 !important;
-    color: rgb(255, 255, 255) !important;
-    line-height: 1 !important;
-    border-radius: 12px !important;
-    transition: box-shadow 200ms ease 0s, background 200ms ease 0s !important;
-    box-shadow: none !important;
-}
-.gr-button-primary:hover{
-    z-index: 14;
-    height: 43px;
-    width: 130px;
-    left: 0px;
-    top: 0px;
-    padding: 0px;
-    cursor: pointer !important;
-    background: none rgb(37, 56, 133) !important;
-    border: none !important;
-    text-align: center !important;
-    font-family: Poppins !important;
-    font-size: 14px !important;
-    font-weight: 500 !important;
-    color: rgb(255, 255, 255) !important;
-    line-height: 1 !important;
-    border-radius: 12px !important;
-    transition: box-shadow 200ms ease 0s, background 200ms ease 0s !important;
-    box-shadow: rgb(0 0 0 / 23%) 0px 1px 7px 0px !important;
-}
-.hover\:bg-orange-50:hover {
-    --tw-bg-opacity: 1 !important;
-    background-color: rgb(229,225,255) !important;
-}
-.to-orange-200 {
-    --tw-gradient-to: rgb(37 56 133 / 37%) !important;
-}
-.from-orange-400 {
-    --tw-gradient-from: rgb(17, 20, 45) !important;
-    --tw-gradient-to: rgb(255 150 51 / 0);
-    --tw-gradient-stops: var(--tw-gradient-from), var(--tw-gradient-to) !important;
-}
-.group-hover\:from-orange-500{
-    --tw-gradient-from:rgb(17, 20, 45) !important; 
-    --tw-gradient-to: rgb(37 56 133 / 37%);
-    --tw-gradient-stops: var(--tw-gradient-from), var(--tw-gradient-to) !important;
-}
-.group:hover .group-hover\:text-orange-500{
-    --tw-text-opacity: 1 !important;
-    color:rgb(37 56 133 / var(--tw-text-opacity)) !important;
-}
-"""
+    testdata = [InputExample(guid=0,text_a=sentence,label=0)]
 
-demo = gr.Interface(fn=fill_mask, inputs=gr.Textbox(lines=1, label="Input"), outputs=gr.Label(label="Output"),title="Fill Mask | Data Science Dojo", theme="light", examples=examples, css=css)
-demo.launch()
+    plm, tokenizer, model_config, WrapperClass = load_plm(type_dic[model_name], model_name)
+
+    promptTemplate = ManualTemplate(
+        text = template,
+        tokenizer = tokenizer,
+    )
+    promptVerbalizer = ManualVerbalizer(
+        classes = classes,
+        label_words = label_words,
+        tokenizer = tokenizer,
+    )
+    test_dataloader = PromptDataLoader(
+        dataset = testdata,
+        tokenizer = tokenizer, 
+        template = promptTemplate, 
+        tokenizer_wrapper_class = WrapperClass,
+        batch_size = 1,
+        max_seq_length = 512,
+    )
+    prompt_model = PromptForClassification(
+        plm=plm,
+        template=promptTemplate, 
+        verbalizer=promptVerbalizer, 
+        freeze_plm=False #whether or not to freeze the pretrained language model
+    )
+    for step, inputs in enumerate(test_dataloader):
+        logits = prompt_model(inputs)
+        
+        
+    return classes[torch.argmax(logits, dim=-1)[0]]
+
+
+demo = gr.Interface(fn=sentiment_analysis,
+                    inputs = [gr.Textbox(placeholder="Enter sentence here.",label="sentence"),
+                            gr.Textbox(placeholder="Your template must have a [SENTENCE] token and a [MASK] token.",label="template"),
+                            gr.Textbox(placeholder="Separate words with Spaces.",label="positive"),
+                            gr.Textbox(placeholder="Separate words with Spaces.",label="neutral"),
+                            gr.Textbox(placeholder="Separate words with Spaces.",label="negative")
+                            ],
+                    outputs="text")
+    
+demo.launch(server_port=8080)