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

tokenizer = None
model = None

def loadmodel():
    tokenizer = AutoTokenizer.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16")
    model = AutoModelForCausalLM.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16", torch_dtype='auto', device_map='auto')
    return tokenizer, model

@spaces.GPU
def generate_text(prompt):
    global tokenizer, model
    if tokenizer is None or model is None:
        tokenizer, model = loadmodel()
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(inputs.input_ids, max_length=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

interface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Meta-Llama-3.1-70B Text Generation",
    description="Enter a prompt and generate text using Meta-Llama-3.1-70B.",
)

interface.launch()