Update base models and datasets

app.py

@@ -6,17 +6,21 @@ import pandas as pd
 import gradio as gr
 import torch
 
+directory = 'av-generation/'
 
-
-def predict(title, type='End2End', num_beams=4):
+def predict(title, type='End2End', base_model='t5-small', dataset='ae-110k', num_beams=3):
     if type == "End2End":
-        model = TransformersAVG('av-generation/t5-small-end2end-ae-110k', use_auth_token=auth_token)
+        model_id = directory + f"{base_model.lower()}-{type.lower()}-{dataset.lower()}"
+        model = TransformersAVG(model_id)
         predictions = model.generate_av_end2end(title, num_beams=num_beams)   
     elif type == "Pipeline":
-        model = TransformersAVG('av-generation/t5-small-ag-ae-110k', model_ve='ksabeh/t5-small-ve-ae-110k', use_auth_token=auth_token)
+        ag_model_id = directory + f"{base_model.lower()}-ag-{dataset.lower()}"
+        ve_model_id = directory + f"{base_model.lower()}-ve-{dataset.lower()}"
+        model = TransformersAVG(ag_model, model_ve=ve_model)
         predictions = model.generate_av_pipeline(title, num_beams=num_beams)
     elif type == 'Multitask':
-        model = TransformersAVG('av-generation/t5-small-mlt-ae-110k', use_auth_token=auth_token)
+        model_id = directory + f"{base_model.lower()}-mlt-{dataset.lower()}"
+        model = TransformersAVG(model_id)
         predictions = model.generate_av_mul(title, num_beams=num_beams)
     else:
         pass
@@ -55,6 +59,10 @@ demo = gr.Interface(
         ),
         gr.Dropdown(
             ["End2End", "Pipeline", "Multitask"], value=["End2End"], multiselect=False, label="AVG Approach", info="Select type of AVG approach."),
+        
+        gr.Radio(["T5-small", "T5-base", "T5-large", "Bart-base", "Bart-large"], value=['T5-small'], label="Base Model", info="Select base model."),
+
+        gr.Radio(["AE-110K", "OA-Mine"], value=["AE-110K"], label="Dataset", info="Select dataset."),
 
         gr.Slider(1, 10, value=4, step=1, label="Number of Beams", info="Degree of exploration at inference.")
     ],