from fastapi import FastAPI, Request
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
import uvicorn

app = FastAPI()

print("Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
    device_map="cpu"
)
print("Base model loaded")

print("Loading NYXA adapter...")
model = PeftModel.from_pretrained(base_model, "ScuraDimensions/NYXA-Mistral-7B")
tokenizer = AutoTokenizer.from_pretrained("ScuraDimensions/NYXA-Mistral-7B")
print("NYXA loaded")

@app.post("/generate")
async def generate(request: Request):
    data = await request.json()
    prompt = data.get("prompt", "")
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=500)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {"response": response}

@app.get("/health")
def health():
    return {"status": "ok"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)