import os
import re
import time
import torch
import spaces
import gradio as gr
from threading import Thread
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer, pipeline

MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"

system_message = {"role":"system", "content":"You are a cybersecurity pentester. You are adept at performing pentesting of networks and machines."}

#pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B")

@spaces.GPU(duration=120)
def process(text):
    messages = []
    messages.append(system_message)
    messages.append({"role": "user", "content": text})
    return pipeline(messages)

def initialize_model():
    """Initialize the model with appropriate configurations"""
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True
    )
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID , trust_remote_code=True)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        device_map="cuda",
        # attn_implementation="flash_attention_2",
        trust_remote_code=True,
        quantization_config=quantization_config

    )
    return model, tokenizer

# Load model directlyxx
#tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Llama-70B")
#model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Llama-70B")
model, tokenizer = initialize_model()
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device="cuda")
demo = gr.Interface(fn=process, inputs="text", outputs="text")
demo.launch()