from language_model import TransformersAVG

import os

import pandas as pd
import gradio as gr
import torch

directory = 'av-generation/'

def predict(title, type='End2End', base_model='t5-small', dataset='ae-110k', num_beams=3):
    if type == "End2End":
        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":
        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_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
    df = pd.DataFrame(predictions, columns=['Attribute', 'Value'])
    return gr.Dataframe(df)

# with gr.Blocks() as demo:
#     gr.Markdown("""
#     # Attribute Value Generation
#     Select Model and AVG Type,  Type into the text box, and click RUN to get Attributes and Values generated by AI.
#     """)

#     title = gr.Textbox(
#             label = "Title",
#             info = "Title of product",
#             lines = 2,
#         )
#     type = gr.Dropdown(
#             ["End2End", "Pipeline", "Multitask"], value=["End2End"], multiselect=False, label="AVG Type", info="Select type of AVG approach.")
    
#     num_beams = gr.Slider(1, 10, value=4, step=1, label="Number of Beams", info="Degree of exploration at inference.")
#     run_btn = gr.Button("Run")
#     output = gr.Dataframe(label="Output Attribute Values")
#     run_btn.click(fn=predict, inputs=[title, type, num_beams], outputs=output, api_name='predict')

# demo.launch()


demo = gr.Interface(
    predict,
    [
        gr.Textbox(
            label = "Title",
            info = "Title of product",
            lines = 2,
        ),
        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.")
    ],
    "dataframe",
    title="Attribute Value Generation",
    examples=[["Women/Girls Tap Dance Shoes Patent Leather Shiny Red /Black/White Tap Shoes for Kids Teacher Practice Performance Shoes T30"], 
              ["2018 Nike Dunk High Premium SB Men's Breathable Hard-wearing Skateboarding Shoes NIKE Sports Sneakers 313171-674"],
              ["New style of high-quality custom LP electric guitar, Abalone Flower inlaid fingerboard electric guitar, maple top, free shipping"],
              ["2Pcs Glasses Side Protection Optical Aye Mate Universal Sideshield Side Shields  Cycling Eyewear"],
              ["Li-Ning 2018 Men Wade Series Jersey Regular Fit 81% Polyester 19% Spandex Breathable Tops Li Ning Sports T-Shirts Tee ATSN149"],
              ["LASPERAL Autumn Winter Fitness Men Running Jackets Coat Sports PU Leather Patchwork Long Sleeve Slim Gym Soccer Baseball Jackets"],
              ["Original New Arrival Authentic NIKE AIR ZOOM VOMERO V12 Men's Breathable Running Shoes Sports Comfortable Sneakers 863762-008"]
              ],
    cache_examples = True
)

demo.launch()