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

app = FastAPI()

# モデルのロード
MODEL_NAME = "Lightblue/DeepSeek-R1-Distill-Qwen-7B-Japanese"
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}