File size: 822 Bytes
2765893
e7692af
c635ea8
 
03fed78
c69f79a
03fed78
 
 
d17dc0b
9fcf154
2765893
 
9fcf154
 
a0fa0c4
 
 
c635ea8
03fed78
2765893
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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")
iface.launch()