Real app?

app.py

@@ -1,25 +1,41 @@
 import gradio as gr
-from transformers import pipeline
+import torch
 import os
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Token aus den Space-Secrets laden
-token = os.getenv("HF_TOKEN")
+MODEL_ID = "patronus-protect/wolf-guard"
+TOKEN = os.getenv("HF_TOKEN")
 
-# Modelle laden
-model_a = pipeline("text-generation", model="username/privat-modell-1", token=token)
-model_b = pipeline("text-generation", model="username/privat-modell-2", token=token)
+# Modell und Tokenizer beim Start der App laden
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=TOKEN)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, token=TOKEN)
 
-def predict(text, model_choice):
-    if model_choice == "Modell 1":
-        return model_a(text)[0]['generated_text']
-    else:
-        return model_b(text)[0]['generated_text']
+def predict(text):
+    label_names = {0: "benign", 1: "attack"}
+    
+    # Text verarbeiten
+    enc = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    
+    with torch.no_grad():
+        logits = model(**enc).logits
+        probs = torch.softmax(logits, dim=-1).squeeze()
+    
+    pred_id = int(probs.argmax())
+    label = label_names.get(pred_id, f"ID {pred_id}")
+    
+    # Rückgabe als Dictionary für das Gradio "Label"-Feld
+    # Zeigt die Wahrscheinlichkeiten beider Klassen an
+    return {label_names[i]: float(probs[i]) for i in range(len(label_names))}
 
-# Gradio Interface
+# Gradio Interface Setup
 demo = gr.Interface(
     fn=predict,
-    inputs=[gr.Textbox(), gr.Dropdown(["Modell 1", "Modell 2"])],
-    outputs="text"
+    inputs=gr.Textbox(label="Eingabe Text", placeholder="Text hier einfügen...", lines=3),
+    outputs=gr.Label(label="Klassifizierung"),
+    title="Sicherheits-Check",
+    description="Dieses Modell erkennt Angriffe in Texten.",
+    allow_flagging="never"
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1 +1,2 @@
-transformers
+transformers
+torch