Update app.py

app.py

@@ -1,30 +1,19 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import DebertaTokenizer, DebertaForSequenceClassification, get_linear_schedule_with_warmup
 import torch
 import re
 from tokenizers import normalizers
 from tokenizers.normalizers import Sequence, Replace, Strip
 from tokenizers import Regex
 
-model1_path = "modernbert.bin"
-model2_path = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed12"
-model3_path = "https://huggingface.co/mihalykiss/modernbert_2/resolve/main/Model_groups_3class_seed22"
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+tokenizer = DebertaTokenizer.from_pretrained('microsoft/deberta-base')
 
-model_1 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
-model_1.load_state_dict(torch.load(model1_path, map_location=device))
-model_1.to(device).eval()
 
-model_2 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
-model_2.load_state_dict(torch.hub.load_state_dict_from_url(model2_path, map_location=device))
+model_2 = DebertaForSequenceClassification.from_pretrained("mihalykiss/best_merged_41_2", num_labels=41)
 model_2.to(device).eval()
 
-model_3 = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
-model_3.load_state_dict(torch.hub.load_state_dict_from_url(model3_path, map_location=device))
-model_3.to(device).eval()
-
 
 label_mapping = {
     0: '13B', 1: '30B', 2: '65B', 3: '7B', 4: 'GLM130B', 5: 'bloom_7b',
@@ -66,15 +55,12 @@ def classify_text(text):
     inputs = tokenizer(cleaned_text, return_tensors="pt", truncation=True, padding=True).to(device)
 
     with torch.no_grad():
-        logits_1 = model_1(**inputs).logits
+
         logits_2 = model_2(**inputs).logits
-        logits_3 = model_3(**inputs).logits
 
-        softmax_1 = torch.softmax(logits_1, dim=1)
         softmax_2 = torch.softmax(logits_2, dim=1)
-        softmax_3 = torch.softmax(logits_3, dim=1)
 
-        averaged_probabilities = (softmax_1 + softmax_2 + softmax_3) / 3
+        averaged_probabilities = softmax_2
         probabilities = averaged_probabilities[0]
 
     ai_probs = probabilities.clone()
@@ -92,7 +78,7 @@ def classify_text(text):
     else:
         result_message = (
             f"**The text is** <span class='highlight-ai'>**{ai_total_prob:.2f}%** likely <b>AI generated</b>.</span>\n\n"
-            f"**Identified AI Model: {ai_argmax_model}**"
+            f"**Identified LLM: {ai_argmax_model}**"
         )
 
     return result_message