aalkaswan commited on
Commit
03fed78
Β·
1 Parent(s): 580d83a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -2,18 +2,23 @@ import gradio as gr
2
  import os
3
  from setfit import SetFitModel
4
 
 
 
 
 
 
 
 
5
 
6
- def greet(name):
7
- return "Hello " + name + "!!"
8
 
9
  # model = SetFitModel.from_pretrained("AISE-TUDelft/java-ownership-classifier")
10
  model = SetFitModel.from_pretrained("dvilasuero/setfit-mini-imdb")
11
 
12
- def classify(text):
13
  if model([text])[0]:
14
  return 'True'
15
  else:
16
  return 'False'
17
 
18
- iface = gr.Interface(fn=classify, inputs=["text", gr.inputs.Dropdown(["gpt2", "gpt-j-6B"])], outputs="text")
19
  iface.launch()
 
2
  import os
3
  from setfit import SetFitModel
4
 
5
+ model_names = ['java-summary', 'java-pointer', 'java-deprecation', 'java-rational', 'java-ownership', 'java-usage', 'java-expand',
6
+ 'pharo-example', 'pharo-key-implementation', 'pharo-responsibilities', 'pharo-collaborators'
7
+ 'python-summary', 'pharo-parameters', 'pharo-usage', 'pharo-development-notes', 'pharo-expand']
8
+
9
+ models = {}
10
+ for model_name in model_names
11
+ models[f'AISE-TUDelft/{model_name}'] = SetFitModel.from_pretrained("dvilasuero/setfit-mini-imdb")#f'AISE-TUDelft/{model_name}'
12
 
 
 
13
 
14
  # model = SetFitModel.from_pretrained("AISE-TUDelft/java-ownership-classifier")
15
  model = SetFitModel.from_pretrained("dvilasuero/setfit-mini-imdb")
16
 
17
+ def classify(text, model):
18
  if model([text])[0]:
19
  return 'True'
20
  else:
21
  return 'False'
22
 
23
+ iface = gr.Interface(fn=classify, inputs=["text", gr.inputs.Dropdown(model_names, label='class')], outputs="text")
24
  iface.launch()