File size: 662 Bytes
2809bee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
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}
|