from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
import torch

app = FastAPI(title="Micro Tier LLM API")

MODELS = {
    "qwen-0.5b": "Qwen/Qwen2.5-0.5B-Instruct",
    "llama-1b": "meta-llama/Llama-3.2-1B-Instruct",
    "qwen-1b": "Qwen/Qwen2.5-1B-Instruct",
    "gemma-1b": "google/gemma-3-1b-it"
}

loaded_models = {}

class GenerateRequest(BaseModel):
    model: str
    prompt: str
    max_new_tokens: int = 256
    temperature: float = 0.7


def get_pipeline(model_name):
    if model_name not in loaded_models:
        loaded_models[model_name] = pipeline(
            "text-generation",
            model=MODELS[model_name],
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto"
        )
    return loaded_models[model_name]


@app.get("/")
def root():
    return {
        "status": "running",
        "models": list(MODELS.keys())
    }


@app.post("/generate")
def generate(req: GenerateRequest):
    if req.model not in MODELS:
        return {"error": "Unknown model"}

    pipe = get_pipeline(req.model)

    output = pipe(
        req.prompt,
        max_new_tokens=req.max_new_tokens,
        temperature=req.temperature,
        do_sample=True,
    )

    return {
        "model": req.model,
        "response": output[0]["generated_text"]
    }