dropdowns

app.py

@@ -2,11 +2,22 @@ import gradio as gr
 import torch
 from transformers import BartForConditionalGeneration, BartTokenizer
 
+# initialize model + tok variables
+model = None
+tok = None
 
-model = BartForConditionalGeneration.from_pretrained("hyechanjun/interview-question-remake")
-tok = BartTokenizer.from_pretrained("hyechanjun/interview-question-remake")
+# pass in Strings of model choice and input text for context
+def genQuestion(model_choice, context):
+    if model_choice=="interview-question-remake":
+        model = BartForConditionalGeneration.from_pretrained("hyechanjun/interview-question-remake")
+        tok = BartTokenizer.from_pretrained("hyechanjun/interview-question-remake")
+    elif model_choice=="interview-length-tagged":
+        model = BartForConditionalGeneration.from_pretrained("hyechanjun/interview-length-tagged")
+        tok = BartTokenizer.from_pretrained("hyechanjun/interview-length-tagged")
+    elif model_choice=="reverse-interview-question":
+        model = BartForConditionalGeneration.from_pretrained("hyechanjun/reverse-interview-question")
+        tok = BartTokenizer.from_pretrained("hyechanjun/reverse-interview-question")
 
-def genQuestion(context):
     inputs = tok(context, return_tensors="pt")
     output = model.generate(inputs["input_ids"], num_beams=4, max_length=64, min_length=9, num_return_sequences=4, diversity_penalty =1.0, num_beam_groups=2)
     final_output = ''
@@ -15,6 +26,6 @@ def genQuestion(context):
         final_output +=  [tok.decode(beam, skip_special_tokens=True, clean_up_tokenization_spaces=False) for beam in output][i] + "\n"
 
     return final_output
-    
-iface = gr.Interface(fn=genQuestion, inputs="text", outputs="text")
+
+iface = gr.Interface(fn=genQuestion, inputs=[gr.inputs.Dropdown(["interview-question-remake", "interview-length-tagged", "reverse-interview-question"]), "text"], outputs="text")
 iface.launch()