from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

MODEL_NAME = "rinna/japanese-gpt-0.5b"  # 500Mの軽量モデル

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, 
    torch_dtype=torch.float16,  # メモリ節約
    low_cpu_mem_usage=True  # メモリ圧縮
)

generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def generate_text(prompt: str, max_length: int = 100):
    return generator(prompt, max_length=max_length)[0]['generated_text']