STACC / app.py
aalkaswan's picture
Update app.py
e7bfb8d
raw
history blame
1.28 kB
import gradio as gr
import os
from setfit import SetFitModel
model_names = ['java-summary', 'java-pointer', 'java-deprecation', 'java-rational', 'java-ownership', 'java-usage', 'java-expand',
'pharo-example', 'pharo-key-implementation', 'pharo-responsibilities', 'pharo-collaborators',
'python-summary', 'pharo-parameters', 'pharo-usage', 'pharo-development-notes', 'pharo-expand']
models = {}
for model_name in model_names:
models[model_name] = SetFitModel.from_pretrained("dvilasuero/setfit-mini-imdb")#f'AISE-TUDelft/{model_name}'
def classify(text, model_name):
if models[model_name]([text])[0]:
return 'True'
else:
return 'False'
iface = gr.Interface(fn=classify,
inputs=["text", gr.inputs.Dropdown(model_names, label='class')],
outputs="text",
title='STACC',
description='''# STACC: Sentence-Transformers Assisted code Comment Classifiers
This app showcases STACC, a collection of SetFit-based Comment classifiers, which were created for the [NLBSE-2023 tool competition](https://nlbse2023.github.io/tools/). More details on the tool itself can be found in the [GitHub repo](https://github.com/AISE-TUDelft/STACC) ''')
iface.launch()