Update app.py

app.py

@@ -7,14 +7,12 @@ import gradio as gr
 model = EncoderDecoderModel.from_pretrained("imamnurby/rob2rand_chen_w_prefix_c_fc")
 tokenizer = RobertaTokenizer.from_pretrained("imamnurby/rob2rand_chen_w_prefix_c_fc")
 
-from transformers import RobertaTokenizer, EncoderDecoderModel
-import pandas as pd
-def generate_preds(desc):
-    # input_desc = desc.lower()
-    # if gen_mode=="Channel":
-    #     input_desc = "GENERATE TRIGGER AND ACTION CHANNEL ONLY <pf> " + input_desc 
-    # elif gen_mode=="Function":
-    #     input_desc = "GENERATE BOTH CHANNEL AND FUNCTION FOR TRIGGER AND ACTION <pf> " + input_desc 
+def generate_preds(desc, gen_mode, num_beams, num_returned_seqs):
+    desc = desc.lower()
+    if gen_mode=="Channel":
+        desc = "GENERATE TRIGGER AND ACTION CHANNEL ONLY <pf> " + desc 
+    elif gen_mode=="Function":
+        desc = "GENERATE BOTH CHANNEL AND FUNCTION FOR TRIGGER AND ACTION <pf> " + desc 
         
     input_ids = tokenizer.encode(desc, return_tensors='pt')
     
@@ -22,8 +20,49 @@ def generate_preds(desc):
     preds = model.generate(
         input_ids, 
         max_length=100, 
-        num_beams=10,
-        num_return_sequences=10,
+        num_beams=num_beams,
+        num_return_sequences=num_returned_seqs,
+        early_stopping=True
+    )
+    
+    output_list = []
+    for item in preds:
+        output_list.append(tokenizer.decode(item, skip_special_tokens=True))
+    
+    if gen_mode=="Channel":
+        trigger = [x.split("<sep>")[0].strip() for x in output_list]
+        trigger_desc = ["dummy" for x in output_list]
+        action = [x.split("<sep>")[1].strip() for x in output_list]
+        action_desc = ["dummy" for x in output_list]
+        df = {"Trigger": trigger,
+              "Trigger Description": trigger_desc,
+              "Action": action,
+              "Action Description": action_desc
+             }
+    elif gen_mode=="Function":
+        trigger = [x.split("<sep>")[1].strip() for x in output_list]
+        trigger_desc = ["dummy" for x in output_list]
+        action = [x.split("<sep>")[3].strip() for x in output_list]
+        action_desc = ["dummy" for x in output_list]
+        df = {"Trigger": trigger,
+              "Trigger Description": trigger_desc,
+              "Action": action,
+              "Action Description": action_desc
+             }
+    return pd.DataFrame(df)
+
+
+def generate_preds_field(desc_field, num_beams_field, num_returned_seqs_field):
+    desc_field = desc_field.lower()
+    
+    input_ids = tokenizer.encode(desc_field, return_tensors='pt')
+    
+    # activate beam search and early_stopping
+    preds = model.generate(
+        input_ids, 
+        max_length=100, 
+        num_beams=num_beams_field,
+        num_return_sequences=num_returned_seqs_field,
         early_stopping=True
     )
     
@@ -31,28 +70,22 @@ def generate_preds(desc):
     for item in preds:
         output_list.append(tokenizer.decode(item, skip_special_tokens=True))
     
-    # if gen_mode=="Channel":
     trigger = [x.split("<sep>")[0].strip() for x in output_list]
-    # trigger_desc = ["dummy" for x in output_list]
+    trigger_desc = ["dummy" for x in output_list]
+    trigger_fields = ["dummy" for x in output_list]
     action = [x.split("<sep>")[1].strip() for x in output_list]
-    # action_desc = ["dummy" for x in output_list]
+    action_desc = ["dummy" for x in output_list]
+    action_fields = ["dummy" for x in output_list]
     df = {"Trigger": trigger,
-          # "Trigger Description": trigger_desc,
+          "Trigger Description": trigger_desc,
+          "Trigger Fields": trigger_fields,
           "Action": action,
-          # "Action Description": action_desc
+          "Action Description": action_desc,
+          "Action Fields": action_fields
          }
-    # elif gen_mode=="Function":
-    #     trigger = [x.split("<sep>")[1].strip() for x in output_list]
-    #     trigger_desc = ["dummy" for x in output_list]
-    #     action = [x.split("<sep>")[3].strip() for x in output_list]
-    #     action_desc = ["dummy" for x in output_list]
-    #     df = {"Trigger": trigger,
-    #           # "Trigger Description": trigger_desc,
-    #           "Action": action,
-    #           # "Action Description": action_desc
-    #          }
     return pd.DataFrame(df)
 
+
 demo = gr.Blocks()
 with demo:
     gr.Markdown("<h1><center>RecipeGen: Automated TAPs Generation Tool</center></h1>")
@@ -69,24 +102,27 @@ with demo:
     with gr.Tabs():
         with gr.TabItem("Channel/Function"):    
             with gr.Column():
-                # gen_mode = gr.Radio(label="Granularity", choices=["Channel", "Function"])
+                gen_mode = gr.Radio(label="Granularity", choices=["Channel", "Function"])
                 desc = gr.Textbox(label="Functionality Description", placeholder="Describe the functionality here")
-                # num_beams = gr.Slider(minimum=2, maximum=500, value=2, step=1, label="Beam Width")
-                # num_returned_seqs = gr.Slider(minimum=2, maximum=500, value=2, step=1, label="#Returned Sequences")
-                # with gr.Row():
+                num_beams = gr.Slider(minimum=2, maximum=500, value=10, step=1, label="Beam Width")
+                num_returned_seqs = gr.Slider(minimum=2, maximum=500, value=10, step=1, label="#Returned Sequences")
                 generate = gr.Button("Generate")
-                # gr.Markdown("<h1><center>Results</center></h1>")
+                
+            with gr.Box():
+                gr.Markdown("<h1><center>Results</center></h1>")
                 results = gr.Dataframe(headers=["Trigger", "Trigger Description", "Action", "Action Description"])
         
-        # with gr.TabItem("Field"):    
-        #     with gr.Column():
-                # desc = gr.Textbox(label="Functionality Description", placeholder="Describe the functionality here")
-                # num_beams = gr.Slider(minimum=2, maximum=500, value=2, step=1, label="Beam Width")
-                # num_returned_seqs = gr.Slider(minimum=2, maximum=500, value=2, step=1, label="#Returned Sequences")
-                # with gr.Row():
-                #     generate_field = gr.Button("Generate")
-                # gr.Markdown("<h1><center>Results</center></h1>")
-                # results_field = gr.Dataframe(headers=["Trigger", "Trigger Description", "Trigger Fields", "Action", "Action Description", "Action Fields"])
+        with gr.TabItem("Field"):    
+            with gr.Column():
+                desc_field = gr.Textbox(label="Functionality Description", placeholder="Describe the functionality here")
+                num_beams_field = gr.Slider(minimum=2, maximum=500, value=10, step=1, label="Beam Width")
+                num_returned_seqs_field = gr.Slider(minimum=2, maximum=500, value=10, step=1, label="#Returned Sequences")
+                generate_field = gr.Button("Generate")
+                
+            with gr.Box():
+                gr.Markdown("<h1><center>Results</center></h1>")
+                results_field = gr.Dataframe(headers=["Trigger", "Trigger Description", "Trigger Fields", "Action", "Action Description", "Action Fields"])
     
-    generate.click(generate_preds, inputs=[desc], outputs=[results])
+    generate.click(generate_preds, inputs=[desc, gen_mode, num_beams, num_returned_seqs], outputs=[results])
+    generate_field.click(generate_preds_field, inputs=[desc_field, num_beams_field, num_returned_seqs_field], outputs=[results_field])
 demo.launch()