UI - theme

app.py

@@ -8,8 +8,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
 model = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", num_labels=41)
 model.load_state_dict(torch.load(model_path, map_location=device))
-model.to(device)
-model.eval()
+model.to(device).eval()
 
 label_mapping = {
     0: '13B', 1: '30B', 2: '65B', 3: '7B', 4: 'GLM130B', 5: 'bloom_7b',
@@ -27,28 +26,34 @@ label_mapping = {
 def classify_text(text):
     if not text.strip():
         return "----"
-    inputs = tokenizer(text, return_tensors="pt", truncation=True)
-    inputs = {key: value.to(device) for key, value in inputs.items()}
-    
+
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+
     with torch.no_grad():
-        outputs = model(**inputs)
-        probabilities = torch.softmax(outputs.logits, dim=1)[0]
-        predicted_class = torch.argmax(probabilities).item()
-        confidence = probabilities[predicted_class].item() * 100
-
-        if predicted_class == 24:
-            prediction_label = f"✅ - The text is <span class='highlight-human'>**{confidence:.2f}%** likely <b>Human written</b>.</span>"
-            model_info = ""  
-        else:
-            prediction_label = f"🤖 - The text is <span class='highlight-ai'>**{confidence:.2f}%** likely <b>AI generated</b>.</span>"
-            model_info = f"**Identified AI Model:** {label_mapping[predicted_class]}"
-
-    result_message = f"**Result:**\n\n{prediction_label}"
-    if model_info:
-        result_message += f"\n\n{model_info}"
-        
+        probabilities = torch.softmax(model(**inputs).logits, dim=1)[0]
+
+    human_prob = probabilities[24].item() * 100
+    ai_probs = probabilities.clone()
+    ai_probs[24] = 0
+    ai_total_prob = ai_probs.sum().item() * 100
+
+    ai_argmax_index = torch.argmax(ai_probs).item()
+    ai_argmax_model = label_mapping[ai_argmax_index]
+
+    if human_prob > ai_total_prob:
+        result_message = (
+            f"✅ - The text is <span class='highlight-human'>**{human_prob:.2f}%** likely <b>Human written</b>.</span>"
+        )
+    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}"
+        )
+
     return result_message
 
+
+
 title = "AI Text Detector"
 
 description = """