Update app.py

app.py

@@ -1,10 +1,23 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import os
+import pandas as pd
+from mitreattack.stix20 import MitreAttackData
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# Chemins des fichiers JSON
+ics_attack_path = 'ics-attack.json'
+enterprise_attack_path = 'enterprise-attack.json'
+
+# Charger les données ATT&CK
+mitre_attack_data = MitreAttackData(enterprise_attack_path)
+
+# Charger les techniques ATT&CK
+techniques = mitre_attack_data.get_techniques(remove_revoked_deprecated=True)
+
+# Convert techniques to a readable string format
+techniques_str = "\n".join([f"{technique['name']} ({mitre_attack_data.get_attack_id(technique['id'])})" for technique in techniques])
+
+client = InferenceClient(model='mistralai/Mixtral-8x7B-Instruct-v0.1')
 
 
 def respond(
@@ -42,22 +55,18 @@ def respond(
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
+
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value=f"""<s>[INST] Given these TTPs: {techniques_str}\n\nfigure out which technique is used in these logs and respond in bullets points and nothing else:\n{log_line}\n[/INST]""", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
 )
 
 
+
 if __name__ == "__main__":
     demo.launch()