app.py

import os

os.environ['HF_HOME'] = '/home/user/.cache/huggingface' 
os.environ['TRANSFORMERS_CACHE'] = os.environ['HF_HOME'] 

from fastapi import FastAPI
from qwen_classifier.predict import predict_single  # Your existing function
import torch
from huggingface_hub import login
from qwen_classifier.model import QwenClassifier
import os

app = FastAPI(title="Qwen Classifier")

@app.on_event("startup")
async def load_model():
    # Warm up GPU
    torch.zeros(1).cuda() 
    # Read HF_TOKEN from Hugging Face Space secrets
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise ValueError("HF_TOKEN not found in environment variables")

    # Authenticate
    login(token=hf_token)
    
    # Load model (will cache in /home/user/.cache/huggingface)
    app.state.model = QwenClassifier.from_pretrained(
        'KeivanR/Qwen2.5-1.5B-Instruct-MLB-clf_lora-1743189446'
    )
    print("Model loaded successfully!")

@app.post("/predict")
async def predict(text: str):
    return predict_single(text, backend="local")