Test update

app.py

@@ -1,40 +1,64 @@
 import gradio as gr
-import torch
 import os
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline
 
-MODEL_ID = "patronus-protect/wolf-guard"
+# Konfiguration
+MODEL_1_ID = "patronus-protect/wolf-guard"
+MODEL_2_ID = "HuggingLil/pii-sensitive-ner-german"
 TOKEN = os.getenv("HF_TOKEN")
+LABEL_NAMES = {0: "benign", 1: "attack"}
 
-# Modell und Tokenizer beim Start der App laden
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=TOKEN)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, token=TOKEN)
+# Pipelines laden (einfachster Weg für verschiedene Modelltypen)
+pipe_cls = pipeline("text-classification", model=MODEL_1_ID, token=TOKEN)
+pipe_ner = pipeline("ner", model=MODEL_2_ID, token=TOKEN, aggregation_strategy="simple")
 
-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}")
+def process_realtime(text):
+    if not text or text.strip() == "":
+        return {}, []
+
+    # Modell 1: Classification (Benign vs Attack)
+    # return_all_scores=True gibt uns die Prozentzahlen für alle Klassen
+    res_cls = pipe_cls(text, top_k=None)
+    cls_output = {item['label']: float(item['score']) for item in res_cls}
+
+    # Modell 2: NER
+    res_ner = pipe_ner(text)
+    # Wir formatieren die NER-Ergebnisse als Liste von Tupeln für die Anzeige
+    # Oder als Dictionary für ein zweites Label-Feld, wenn du nur die Entities willst
+    ner_output = {f"{item['entity_group']}: {item['word']}": float(item['score']) for item in res_ner}
+
+    return cls_output, ner_output
+
+# Custom CSS für Framer-Integration
+css = """
+footer {display: none !important;}
+.gradio-container {background: transparent !important; padding: 0 !important;}
+"""
+
+with gr.Blocks(css=css) as demo:
+    input_text = gr.Textbox(
+        label="Eingabe Analyse", 
+        placeholder="Text eingeben...", 
+        lines=3,
+        interactive=True
+    )
     
-    # 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 Setup
-demo = gr.Interface(
-    fn=predict,
-    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."
-)
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Modell 1: Klassifizierung")
+            output_cls = gr.Label(label="Wahrscheinlichkeit")
+            
+        with gr.Column():
+            gr.Markdown("### Modell 2: NER")
+            output_ner = gr.Label(label="Erkannte Entitäten")
+
+    # Realtime Trigger
+    input_text.change(
+        fn=process_realtime, 
+        inputs=input_text, 
+        outputs=[output_cls, output_ner],
+        show_progress="hidden"
+    )
 
 if __name__ == "__main__":
     demo.launch()