from fastapi import FastAPI
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import torch

app = FastAPI()

model_name = "Jaswant1801/qwen_nutrisync"

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

# Load config and remove quantization
config = AutoConfig.from_pretrained(model_name)
config.quantization_config = None

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    config=config,
    trust_remote_code=True,
    device_map="cpu",
    dtype=torch.float32
)

@app.get("/")
def home():
    return {"message": "NutriSync API running"}

@app.post("/generate")
def generate(prompt: str):

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

    outputs = model.generate(
        **inputs,
        max_new_tokens=150
    )

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

    return {"response": response}