Update app.py

app.py

@@ -1,12 +1,17 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
-MODEL_NAME = "Lightblue/DeepSeek-R1-Distill-Qwen-7B-Japanese"
+MODEL_NAME = "rinna/japanese-gpt-0.5b"  # 500Mの軽量モデル
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.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)
 
-prompt = "こんにちは、これはテストです。"
-result = generator(prompt, max_length=100)[0]['generated_text']
-print(result)
+def generate_text(prompt: str, max_length: int = 100):
+    return generator(prompt, max_length=max_length)[0]['generated_text']
+