import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import spaces

peft_model_id = "rootxhacker/CodeAstra-7B"
config = PeftConfig.from_pretrained(peft_model_id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    return_dict=True,
    load_in_4bit=True,
    device_map="auto"  # This will automatically handle device placement
    )

tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

model = PeftModel.from_pretrained(model, peft_model_id)


@spaces.GPU(duration=200)
def get_completion(query, model, tokenizer):
    inputs = tokenizer(query, return_tensors="pt").to(device)  # Move inputs to the same device as the model
    outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

@spaces.GPU(duration=200)
def code_review(code_to_analyze):
    query = f"As a code review expert, your role will be to carefully examine the code for potential security flaws and provide guidance on secure coding practices. This may include identifying common coding mistakes that could lead to vulnerabilities, suggesting ways to improve the code's overall security, and recommending tools or techniques that can be used to detect and prevent potential threats. Your expertise in security will be particularly valuable in ensuring that any code developed meets the highest security standard:\n{code_to_analyze}"
    result = get_completion(query, model, tokenizer)
    return result

# Create Gradio interface
iface = gr.Interface(
    fn=code_review,
    inputs=gr.Textbox(lines=10, label="Enter code to analyze"),
    outputs=gr.Textbox(label="Code Review Result"),
    title="Code Review Expert",
    description="This tool analyzes code for potential security flaws and provides guidance on secure coding practices."
)

# Launch the Gradio app with a public link
iface.launch()