src/open_r1/main.py

from fastapi import FastAPI
from pydantic import BaseModel
from typing import Optional

from open_r1.evaluate import TASKS_TABLE
from open_r1.generate import build_distilabel_pipeline

app = FastAPI()


class GenerateRequest(BaseModel):
    model: str
    prompt_column: Optional[str] = None
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    max_new_tokens: int = 8192
    num_generations: int = 1


@app.get("/")
def root():
    return {"message": "Welcome to DeepSeek-R1!"}


@app.get("/tasks")
def get_tasks():
    """Endpoint to fetch available tasks"""
    return {"tasks": [task.name for task in TASKS_TABLE]}


@app.post("/generate")
def generate_text(request: GenerateRequest):
    """Endpoint for text generation"""
    pipeline = build_distilabel_pipeline(
        model=request.model,
        prompt_column=request.prompt_column,
        temperature=request.temperature,
        top_p=request.top_p,
        max_new_tokens=request.max_new_tokens,
        num_generations=request.num_generations,
    )
    return {"message": "Pipeline successfully built. Use it to run your dataset."}