import gradio as gr
from transformers import pipeline


pipe = pipeline(
    "text-generation",
    model="niranjan2777/cybersec-qwen",
    tokenizer="niranjan2777/cybersec-qwen",
    device_map="cpu",
    torch_dtype="auto",
)


def cybersec_assistant(user_input):
    messages = [
        {"role": "system", "content": "You are a cybersecurity expert. Give clear, structured answers."},
        {"role": "user", "content": user_input},
    ]

    output = pipe(
        messages,
        max_new_tokens=200,
        temperature=0.7
    )

    return output[0]["generated_text"][-1]["content"]


iface = gr.Interface(
    fn=cybersec_assistant,
    inputs=gr.Textbox(lines=3, placeholder="Ask a cybersecurity question..."),
    outputs="text",
    title=" CyberSec Assistant_ Niranjan Labs",
    description="Ask anything about cybersecurity.."
)

iface.launch()