from fastapi import FastAPI
from pydantic import BaseModel

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM
)

import torch

app = FastAPI()

MODEL_ID = "himalaya-ai/himalayagpt-0.5b"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.float32
)

class Request(BaseModel):
    prompt: str
    max_tokens: int = 100

@app.get("/")
def home():
    return {"status": "running"}

@app.post("/generate")
def generate(req: Request):

    inputs = tokenizer(
        req.prompt,
        return_tensors="pt"
    )

    outputs = model.generate(
        **inputs,
        max_new_tokens=req.max_tokens,
        temperature=0.7,
        do_sample=True
    )

    response = tokenizer.decode(
        outputs[0],
        skip_special_tokens=True
    )

    return {
        "response": response
    }