Qifan Zhang
add dataframe in frontend
b49f004
raw
history blame
2.48 kB
from io import StringIO
from typing import Optional
import gradio as gr
import pandas as pd
from utils import pipeline
from utils.models import list_models
def read_data(filepath: str) -> Optional[pd.DataFrame]:
if filepath.endswith('.xlsx'):
df = pd.read_csv(filepath)
elif filepath.endswith('.csv'):
df = pd.read_csv(filepath)
else:
raise Exception('File type not supported')
return df
def process(task_name: str,
model_name: str,
text: str,
file=None,
):
try:
# load file
if file:
df = read_data(file.name)
elif text:
string_io = StringIO(text)
df = pd.read_csv(string_io)
assert len(df) >= 1, 'No input data'
else:
raise Exception('No input data')
# process
if task_name == 'Originality':
df = pipeline.p0_originality(df, model_name)
elif task_name == 'Flexibility':
df = pipeline.p1_flexibility(df, model_name)
else:
raise Exception('Task not supported')
# save
path = 'output.csv'
df.to_csv(path, index=False, encoding='utf-8-sig')
return None, df.iloc[:10], path
except Exception as e:
return {'Error': e}, None, None
instructions = 'Please upload a file or paste the text below. '
task_name_dropdown = gr.components.Dropdown(
label='Task Name',
value='Originality',
choices=['Originality', 'Flexibility']
)
model_name_input = gr.components.Textbox(
value='paraphrase-multilingual-MiniLM-L12-v2',
lines=1,
type='text'
)
model_name_dropdown = gr.components.Dropdown(
label='Model Name',
value=list_models[0],
choices=list_models
)
text_input = gr.components.Textbox(
value='id,prompt,response\n',
lines=10,
type='text'
)
text_output = gr.components.Textbox(
label='Output',
type='text'
)
dataframe_output = gr.components.Dataframe(
label='DataFrame'
)
description = open('description.txt', 'r').read()
file_output = gr.components.File(label='Output File',
file_count='single',
file_types=['', '.', '.csv', '.xls', '.xlsx'])
app = gr.Interface(
fn=process,
inputs=[task_name_dropdown, model_name_dropdown, text_input, 'file'],
outputs=[text_output, dataframe_output, file_output],
description=description
)
app.launch()