app.py

import os


from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Load model & tokenizer once at startup
MODEL_ID = "EQuIP-Queries/EQuIP_3B"
# Specify cache_dir just in case
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)

# 2. Initialize FastAPI
app = FastAPI()

# 3. Define request schema
class GenerateRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 50

# 4. Inference endpoint
@app.post("/generate")
async def generate(req: GenerateRequest):
    inputs = tokenizer(req.prompt, return_tensors="pt")
    ids = model.generate(**inputs, max_new_tokens=req.max_new_tokens)
    text = tokenizer.decode(ids[0], skip_special_tokens=True)
    return {"generated_text": text}