Update app.py

app.py

@@ -8,14 +8,11 @@ model_names = ['java-summary', 'java-pointer', 'java-deprecation', 'java-rationa
 
 models = {}
 for model_name in model_names:
-    models[f'AISE-TUDelft/{model_name}'] = SetFitModel.from_pretrained("dvilasuero/setfit-mini-imdb")#f'AISE-TUDelft/{model_name}'
+    models[model_name] = SetFitModel.from_pretrained("dvilasuero/setfit-mini-imdb")#f'AISE-TUDelft/{model_name}'
 
 
-# model = SetFitModel.from_pretrained("AISE-TUDelft/java-ownership-classifier")
-model = SetFitModel.from_pretrained("dvilasuero/setfit-mini-imdb")
-
-def classify(text, model):
-    if model([text])[0]:
+def classify(text, model_name):
+    if models[model_name]([text])[0]:
         return 'True'
     else:
         return 'False'