File size: 571 Bytes
2809bee c423e28 2809bee c423e28 c815242 c423e28 2809bee c423e28 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
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']
|