tcyang/TransDis-CreativityAutoAssessment-V2

1import traceback2from io import StringIO3from typing import Optional4 5import gradio as gr6import pandas as pd7from loguru import logger8 9from utils import pipeline10from utils.models import list_models11 12 13def read_data(filepath: str) -> Optional[pd.DataFrame]:14    if filepath.endswith('.xlsx'):15        df = pd.read_excel(filepath)16    elif filepath.endswith('.csv'):17        df = pd.read_csv(filepath)18    else:19        raise Exception('File type not supported')20    return df21 22 23def process(24        task_name: str,25        model_name: str,26        pooling: str,27        text: str,28        file=None,29) -> (None, pd.DataFrame, str):30    try:31        logger.info(f'Processing {task_name} with {model_name} and {pooling}')32        # load file33        if file:34            df = read_data(file.name)35        elif text:36            string_io = StringIO(text)37            df = pd.read_csv(string_io)38            assert len(df) >= 1, 'No input data'39        else:40            raise Exception('No input data')41 42        # check43        if len(df) > 10000:44            raise Exception('Data exceeds 10,000 rows')45 46        # process47        if task_name == 'Originality':48            df = pipeline.p0_originality(df, model_name, pooling)49        elif task_name == 'Flexibility':50            df = pipeline.p1_flexibility(df, model_name, pooling)51        else:52            raise Exception('Task not supported')53 54        # save55        path = 'output.csv'56        df.to_csv(path, index=False, encoding='utf-8-sig')57        return None, df.iloc[:10], path58 59    except:60        error = traceback.format_exc()61        logger.warning({62            'error': error,63            'task_name': task_name,64            'model_name': model_name,65            'pooling': pooling,66            'text': text,67            'file': file,68        })69        return {'Info': 'Something wrong', 'Error': traceback.format_exc()}, None, None70 71 72# input73task_name_dropdown = gr.components.Dropdown(74    label='Task Name',75    value='Originality',76    choices=['Originality', 'Flexibility']77)78model_name_dropdown = gr.components.Dropdown(79    label='Model Name',80    value=list_models[0],81    choices=list_models82)83pooling_dropdown = gr.components.Dropdown(84    label='Pooling',85    value='mean',86    choices=['mean', 'cls']87)88text_input = gr.components.Textbox(89    value=open('data/example_xlm.csv', 'r').read(),90    lines=10,91)92file_input = gr.components.File(label='Input File', file_types=['.csv', '.xlsx'])93 94# output95text_output = gr.components.Textbox(label='Output')96dataframe_output = gr.components.Dataframe(label='DataFrame')97file_output = gr.components.File(label='Output File', file_types=['.csv', '.xlsx'])98 99app = gr.Interface(100    fn=process,101    inputs=[task_name_dropdown, model_name_dropdown, pooling_dropdown, text_input, file_input],102    outputs=[text_output, dataframe_output, file_output],103    description=open('data/description.txt', 'r').read(),104    title='TransDis-CreativityAutoAssessment',105    concurrency_limit=1,106)107app.launch(max_threads=1)108