import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load the model and tokenizer
model_path = 'LLM4Binary/llm4decompile-1.3b-v1.5'  # V1.5 Model
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).cuda()

# Define the inference function
def generate_response(input_text, temperature, top_k, top_p):
    before = f"# This is the assembly code:\n"#prompt
    after = "\n# What is the source code?\n"#prompt
    input_func = before+input_text.strip()+after
    inputs = tokenizer(input_func, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_length=512,  # Adjust this if needed
        do_sample=True,
        top_k=int(top_k),
        top_p=float(top_p),
        temperature=float(temperature)
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Create a Gradio interface with sliders
interface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=5, placeholder="Enter your input text here...", label="Input Text"),
        gr.Slider(0.1, 2.0, value=0.0, step=0.1, label="Temperature"),
        gr.Slider(1, 100, value=10, step=1, label="Top-k"),
        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
    ],
    outputs=gr.Textbox(label="Generated Response"),
    title="LLM4Binary Interactive Demo",
    description="Adjust the sliders for temperature, top-k, and top-p to customize the model's response."
)

# Launch the Gradio app
interface.launch()