Update app

app.py

@@ -14,36 +14,62 @@ first_model = AutoModelForSequenceClassification.from_pretrained(first_model_pat
 second_model_path = "models/second_model"
 second_model = AutoModelForSequenceClassification.from_pretrained(second_model_path)
 
-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")
+f_30_40_model_path = "models/FEMALE_30_40model"
+f_30_40_model = AutoModelForSequenceClassification.from_pretrained(f_30_40_model_path)
 
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=1)
+f_40_55_model_path = "models/FEMALE_40_55model"
+f_40_55_model = AutoModelForSequenceClassification.from_pretrained(f_40_55_model_path)
 
-    predicted_class = torch.argmax(probabilities, dim=1).item()
-    class_labels = model.config.id2label
-    predicted_label = class_labels[predicted_class]
-    probabilities = probabilities[0].tolist()
-
-    categories = ['Kvinner 30-40', 'Kvinner 40-55', 'Menn 30-40', 'Menn 40-55']
+m_30_40_model_path = "models/MALE_30_40model"
+m_30_40_model = AutoModelForSequenceClassification.from_pretrained(m_30_40_model_path)
 
-    #category_probabilities = list(zip(categories, probabilities))
+m_40_55_model_path = "models/MALE_40_55model"
+m_40_55_model = AutoModelForSequenceClassification.from_pretrained(m_40_55_model_path)
 
-    #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])
+def classify_text(test_text, selected_model):
 
-    return dict(zip(categories, map(float,probabilities)))
+    categories = ['Kvinner 30-40', 'Kvinner 40-55', 'Menn 30-40', 'Menn 40-55']
+    
+    if selected_model in ('Model 1', 'Model 2'):
+        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():
+            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()
+
+        return dict(zip(categories, map(float,probabilities)))
+    elif selected_model == 'Model 3':
+        models = [f_30_40_model, f_40_55_model, m_30_40_model, m_40_55_model]
+        performance_labels = []
+        inputs = tokenizer(test_text, return_tensors="pt")
+
+        for model in models:
+            with torch.no_grad():
+                outputs = model(**inputs)
+                logits = outputs.logits
+                probabilities = torch.softmax(logits, dim=1)
+            
+            predicted_class = torch.argmax(probabilities, dim=1).item()
+            performance_labels = ['Lite god', 'Nokså god', 'God']
+            predicted_performance = performance_labels[predicted_class]
+            
+            class_labels = model.config.id2label
+            predicted_label = class_labels[predicted_class]
+            performance_labels.append(predicted_label)
+
+        return dict(zip(categories, map(float,performance_labels)))
 
 # Cell
 label = gr.outputs.Label()
@@ -52,7 +78,10 @@ app_title = "Target group classifier"
 
 examples = [["Moren leter etter sønnen i et ihjelbombet leilighetskompleks.", 'Model 1'],
             ["Fotballstadion tok fyr i helgen", 'Model 2'],
-            ["De første månedene av krigen gikk så som så. Nå har Putin skiftet strategi.", 'Model 1']
+            ["De første månedene av krigen gikk så som så. Nå har Putin skiftet strategi.", 'Model 1'],
+            ["Title: Disse hadde størst formue i 2022, Text lead: Laksearvingen Gustav Magnar Witzøe økte formuen med nesten 7 milliarder i fjor, og troner nok en gang øverst på listen over Norges rikeste.", "Model 3"],
+            ["Title: Dette er de mest populære navnene i 2022, Text lead: Navnetoppen for 2022 er klar. Se hvilke navn som er mest populære i din kommune.", "Model 3"],
+            ["Title: 2023 er det varmeste året noen gang registrert, Text lead: En ny rapport viser at 2023 er det varmeste året registrert siden man startet målingene. Klimaforsker kaller tallene urovekkende.", "Model 3"]
 ]
-intf = gr.Interface(fn=classify_text, inputs=["text", gr.Dropdown(['Model 1', 'Model 2'])], outputs=label, examples=examples, title=app_title)
+intf = gr.Interface(fn=classify_text, inputs=["text", gr.Dropdown(['Model 1', 'Model 2', 'Model 3'])], outputs=label, examples=examples, title=app_title)
 intf.launch(inline=False)