from transformers import AutoTokenizer, AutoModelForCausalLM

def model_fn(model_dir):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(model_dir)
    return {"model": model, "tokenizer": tokenizer}

def predict_fn(data, model_dict):
    model = model_dict["model"]
    tokenizer = model_dict["tokenizer"]

    inputs = tokenizer(data["inputs"], return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_length=128)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def output_fn(prediction, accept):
    return {"generated_text": prediction}