from fastapi import FastAPI
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

app = FastAPI()

# 🚀 軽量モデルに変更（7B → 3.6B）
MODEL_NAME = "rinna/japanese-gpt-neox-3.6b"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

@app.get("/")
def root():
    return {"message": "DeepSeek 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']
    return {"response": result}