| 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'] | |