|
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} |
|
|
