app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
import pandas as pd
from mitreattack.stix20 import MitreAttackData

# 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(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

"""
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=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 if there is no TTPs detected just say that there is not TTP detected: [/INST]""", label="System message"),
            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)"),
    ],
)



if __name__ == "__main__":
    demo.launch()