from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("bartowski/Qwen2.5-Coder-32B-Instruct-abliterated-GGUF")
model = AutoModelForCausalLM.from_pretrained(
    "bartowski/Qwen2.5-Coder-32B-Instruct-abliterated-GGUF",
    device_map="auto",
    torch_dtype="auto",
    resume_download=True  # Enable resumable downloads
)

# Define a function for generating text
def generate_text(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_length=200)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create a Gradio interface
interface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Qwen 2.5 Coder 32B Text Generator",
    description="Enter a prompt to generate text using the Qwen2.5-Coder-32B-Instruct-abliterated-GGUF model."
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()