Update app with second model

app.py

@@ -1,6 +1,5 @@
 __all__ = ['is_flower', 'learn', 'classify_image', 'categories', 'image', 'label', 'examples', 'intf']
 
-from fastai.vision.all import *
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
 import torch
@@ -9,11 +8,21 @@ import torch
 model_name = "NbAiLab/nb-bert-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-model_path = "models/first_model"
-model = AutoModelForSequenceClassification.from_pretrained(model_path)
+first_model_path = "models/first_model"
+first_model = AutoModelForSequenceClassification.from_pretrained(first_model_path)
 
+second_model_path = "models/second_model"
+second_model = AutoModelForSequenceClassification.from_pretrained(second_model_path)
 
-def classify_text(test_text):
+def classify_text(test_text, selected_model):
+
+    if selected_model == 'Model 1':
+        model = first_model
+    elif selected_model == 'Model 2':
+        model = second_model
+    else:
+        raise ValueError("Invalid model selection")
+    
     inputs = tokenizer(test_text, return_tensors="pt")
 
     with torch.no_grad():
@@ -28,9 +37,9 @@ def classify_text(test_text):
 
     categories = ['Kvinner 30-40', 'Kvinner 40-55', 'Menn 30-40', 'Menn 40-55']
 
-    category_probabilities = list(zip(categories, probabilities))
+    #category_probabilities = list(zip(categories, probabilities))
 
-    max_category = max(category_probabilities, key=lambda x: x[1])
+    #max_category = max(category_probabilities, key=lambda x: x[1])
 
     #print('The model predicts that this text lead would have a majority of readers in the target group', max_category[0])
 
@@ -42,8 +51,8 @@ categories = ('Kvinner 30-40', 'Kvinner 40-55', 'Menn 30-40', 'Menn 40-55')
 app_title = "Target group classifier"
 
 examples = ["Moren leter etter sønnen i et ihjelbombet leilighetskompleks.",
-            "De første månedene av krigen gikk så som så. Nå har Putin skiftet strategi."
+            "De første månedene av krigen gikk så som så. Nå har Putin skiftet strategi.",
             "Fotballstadion tok fyr i helgen"
 ]
-intf = gr.Interface(fn=classify_text, inputs="text", outputs=label, examples=examples, title=app_title)
-intf.launch(inline=False)
+intf = gr.Interface(fn=classify_text, inputs=["text", gr.DropDown(['Model 1', 'Model 2'])], outputs=label, examples=examples, title=app_title)
+intf.launch(inline=False)