Change app.py

app.py

@@ -1,27 +1,49 @@
 __all__ = ['is_flower', 'learn', 'classify_image', 'categories', 'image', 'label', 'examples', 'intf']
 
-# Cell
 from fastai.vision.all import *
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
+import torch
 
-def is_cat(x): return x[0].isupper()
 
-# Cell
-learn = load_learner('models/plant_or_flower_model.pkl')
+model_name = "NbAiLab/nb-bert-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Cell
-categories = ('Flower', 'Plant')
+model_path = "/models/first_model"
+model = AutoModelForSequenceClassification.from_pretrained(model_path)
 
-def classify_image(img):
-    pred,idx,probs = learn.predict(img)
-    return dict(zip(categories, map(float,probs)))
 
-# Cell
-image = gr.inputs.Image(shape=(192, 192))
-label = gr.outputs.Label()
-examples = ['images/flower.jpg', 'images/plant.jpg', 'images/dunno.jpg']
+def classify_text(test_text):
+    inputs = tokenizer(test_text, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        probabilities = torch.softmax(logits, dim=1)
+
+    predicted_class = torch.argmax(probabilities, dim=1).item()
+    class_labels = model.config.id2label
+    predicted_label = class_labels[predicted_class]
+    probabilities = probabilities[0].tolist()
 
-app_title = "Plant or Flower Classifier"
+    categories = ['Kvinner 30-40', 'Kvinner 40-55', 'Menn 30-40', 'Menn 40-55']
 
-intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples, title=app_title)
+    category_probabilities = list(zip(categories, probabilities))
+
+    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])
+
+    return dict(zip(categories, map(float,probabilities)))
+
+# Cell
+label = gr.outputs.Label()
+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."
+            "Fotballstadion tok fyr i helgen"
+]
+intf = gr.Interface(fn=classify_text, inputs="text", outputs=label, examples=examples, title=app_title)
 intf.launch(inline=False)