import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model and tokenizer
model_name = "huihui-ai/Llama-3.2-3B-Instruct-abliterated"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    low_cpu_mem_usage=True
)

# Define the text generation function
def generate_text(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(inputs["input_ids"], max_length=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=5, placeholder="Enter your prompt here..."),
    outputs="text",
    title="Llama 3.2 3B Instruct Abliterated",
    description="An uncensored language model. Enter your prompt to receive a response."
)

if __name__ == "__main__":
    iface.launch()