Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 import os
 from mitreattack.stix20 import MitreAttackData
-from descriptions import descriptions
+from descriptions import descriptions  # Assurez-vous que descriptions.py est dans le même répertoire
 
 # Chemins des fichiers JSON
 ics_attack_path = 'ics-attack.json'
@@ -19,6 +19,9 @@ techniques_str = "\n".join([f"{technique['name']} ({mitre_attack_data.get_attack
 
 client = InferenceClient(model='mistralai/Mixtral-8x7B-Instruct-v0.1')
 
+def generate_system_message(log_input):
+    description_output = descriptions(log_input)
+    return f"""<s>[INST] Given these TTPs: {techniques_str}\n\n and here are the descriptions: {description_output}\n\nFigure out which technique is used in these logs and respond in bullet points and nothing else.[/INST]"""
 
 def respond(
     message,
@@ -53,16 +56,32 @@ def respond(
         response += token
         yield response
 
-
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value=f"""<s>[INST] Given these TTPs: {techniques_str}\n\n and here are {descriptions}\n\nfigure out which technique is used in these logs and respond in bullets points and nothing else[/INST]""", label="System message"),
+        gr.Textbox(label="Log Input", placeholder="Enter log here..."),
+        gr.Textbox(label="System message", value="", interactive=False),
         gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.1, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
+    title="TTP Detection Interface",
+    description="Enter logs to detect TTPs using the model.",
 )
 
+# Update system message based on log input
+def update_system_message(log_input):
+    return generate_system_message(log_input)
+
+demo.additional_inputs[1].update(value=gr.update(value=generate_system_message("")))
+
+# Event handler for log input
+def on_log_input_change(log_input):
+    system_message = generate_system_message(log_input)
+    demo.additional_inputs[1].update(value=gr.update(value=system_message))
+
+# Bind event handler
+demo.additional_inputs[0].change(on_log_input_change)
+
 if __name__ == "__main__":
     demo.launch()