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

model_name = "models/unsloth/Meta-Llama-3.1-8B-bnb-4bit"

@spaces.GPU(duration=180)
def load_model():
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

@spaces.GPU
def generate_text(prompt, model, tokenizer):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def gradio_interface():
    model, tokenizer = load_model()
    
    def wrapped_generate(prompt):
        return generate_text(prompt, model, tokenizer)
    
    iface = gr.Interface(
        fn=wrapped_generate,
        inputs="text",
        outputs="text",
        title="Meta-Llama 3.1 8B Text Generation"
    )
    return iface

if __name__ == "__main__":
    demo = gradio_interface()
    demo.launch()