import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
finetuned_model = AutoModelForCausalLM.from_pretrained("AquilaX-AI/security_assistant_2")
tokenizer = AutoTokenizer.from_pretrained("AquilaX-AI/security_assistant")

finetuned_model.to(device)

prompt = """<|im_start|>system
You are a helpful AI assistant named Securitron<|im_end|>
<|im_start|>user
cwe_id:CWE-20
cwe_name:Improper Input Validation
affected_line:Pattern Undefined (v3)
partial_code:example: c4d5ea2f-81a2-4a05-bcd3-202126ae21df
        name:
          type: string
          example: Toolbox
        serial:
file_name:itemit_openapi.yaml
status:True Positive
reason: There is no pattern property that could lead to insufficient input validation.
remediation_action: Always define a pattern to ensure strict input validation.

How to fix this?<|im_end|>
<|im_start|>assistant
"""

s = time.time()

encodeds = tokenizer(prompt, return_tensors="pt",truncation=True).input_ids.to(device)
text_streamer = TextStreamer(tokenizer, skip_prompt = True)

# Increase max_new_tokens if needed
response = finetuned_model.generate(
        input_ids=encodeds,
        streamer=text_streamer,
        max_new_tokens=512,
        use_cache=True,
        pad_token_id=151645,
        eos_token_id=151645,
        num_return_sequences=1
    )
e = time.time()
print(f'time taken:{e-s}')