Update app.py

app.py

@@ -1,20 +1,29 @@
 from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
-app = FastAPI()
+MODEL_NAME = "EQUES/TinyDeepSeek-1.5B"
+
+# メモリ最適化を適用
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME, 
+    torch_dtype=torch.float16,  # メモリ節約のため16bit
+    device_map="auto",  # CPUメモリへ分割割り当て
+    low_cpu_mem_usage=True  # 初期化時のメモリ削減
+)
 
-# 🚀 軽量な日本語モデルを選択（7B → 3.6B）
-MODEL_NAME = "rinna/japanese-gpt-neox-3.6b"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+app = FastAPI()
+
 @app.get("/")
 def root():
-    return {"message": "DeepSeek API is running!"}
+    return {"message": "TinyDeepSeek API is running!"}
 
 @app.get("/generate")
 def generate(prompt: str, max_length: int = 100):
-    result = generator(prompt, max_length=max_length, do_sample=True)[0]['generated_text']
+    result = generator(prompt, max_length=max_length)[0]['generated_text']
     return {"response": result}
 
+