Improve inputs

app.py

@@ -13,13 +13,7 @@ def reduce_mean(value, mask, axis=None):
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
-max_input_len = 256
-max_output_len = 32
-m = 10
-top_p = 0.5
-
 class InteractiveRainier:
-
     def __init__(self):
         self.tokenizer = transformers.AutoTokenizer.from_pretrained('allenai/unifiedqa-t5-large')
         self.rainier_model = transformers.AutoModelForSeq2SeqLM.from_pretrained('liujch1998/rainier-large').to(device)
@@ -46,7 +40,7 @@ class InteractiveRainier:
         choices.append(choice)
         return choices
 
-    def run(self, question):
+    def run(self, question, max_input_len, max_output_len, m, top_p):
         tokenized = self.tokenizer(question, return_tensors='pt', padding='max_length', truncation='longest_first', max_length=max_input_len).to(device) # (1, L)
         knowledges_ids = self.rainier_model.generate(
             input_ids=tokenized.input_ids,
@@ -107,8 +101,8 @@ class InteractiveRainier:
 
 rainier = InteractiveRainier()
 
-def predict(question, choices):
-    result = rainier.run(f'{question} \\n {choices}')
+def predict(question, kg_model, qa_model, max_input_len, max_output_len, m, top_p):
+    result = rainier.run(question, max_input_len, max_output_len, m, top_p)
     output = ''
     output += f'QA model answer without knowledge: {result["knowless_pred"]}\n'
     output += f'QA model answer with knowledge: {result["knowful_pred"]}\n'
@@ -120,13 +114,29 @@ def predict(question, choices):
     output += f'Knowledge selected to make the prediction: {result["selected_knowledge"]}\n'
     return output
 
-input_question = gr.inputs.Textbox(label='Question:')
-input_choices = gr.inputs.Textbox(label='Choices:')
+input_question = gr.inputs.Dropdown(
+    choices=[
+        'If the mass of an object gets bigger what will happen to the amount of matter contained within it? \\n (A) gets bigger (B) gets smaller',
+        'What would vinyl be an odd thing to replace? \\n (A) pants (B) record albums (C) record store (D) cheese (E) wallpaper',
+        'Some pelycosaurs gave rise to reptile ancestral to \\n (A) lamphreys (B) angiosperm (C) mammals (D) paramecium (E) animals (F) protozoa (G) arachnids (H) backbones',
+        'Sydney rubbed Addison’s head because she had a horrible headache. What will happen to Sydney? \\n (A) drift to sleep (B) receive thanks (C) be reprimanded',
+        'Adam always spent all of the free time watching Tv unlike Hunter who volunteered, due to _ being lazy. \\n (A) Adam (B) Hunter',
+        'Causes bad breath and frightens blood-suckers \\n (A) tuna (B) iron (C) trash (D) garlic (E) pubs',
+    ],
+    label='Question:',
+    info='A multiple-choice commonsense question. Please follow the UnifiedQA input format: "{question} \\n (A) ... (B) ... (C) ..."',
+)
+input_kg_model = gr.inputs.Textbox(label='Knowledge generation model:', value='liujch1998/rainier-large', interactive=False)
+input_qa_model = gr.inputs.Textbox(label='QA model:', value='allenai/unifiedqa-t5-large', interactive=False)
+input_max_input_len = gr.inputs.Number(label='Max question length:', value=256, precision=0)
+input_max_output_len = gr.inputs.Number(label='Max knowledge length:', value=32, precision=0)
+input_m = gr.inputs.Slider(label='Number of generated knowledges:', value=10, mininum=1, maximum=20, step=1)
+input_top_p = gr.inputs.Slider(label='Top_p for knowledge generation:', value=0.5, mininum=0.0, maximum=1.0, step=0.05)
 output_text = gr.outputs.Textbox(label='Output')
 
 gr.Interface(
     fn=predict,
-    inputs=[input_question, input_choices],
+    inputs=[input_question, input_kg_model, input_qa_model, input_max_input_len, input_max_output_len, input_m, input_top_p],
     outputs=output_text,
     title="Rainier",
 ).launch()