# Base image with Python and CUDA for GPU support
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04

# Install system dependencies
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
    python3.10 \
    python3-pip \
    python3.10-venv \
    git \
    libgl1 \
    libglib2.0-0 && \
    rm -rf /var/lib/apt/lists/*

# Create and activate virtual environment
RUN python3.10 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# Install Python dependencies first (for better caching)
COPY requirements.txt .
COPY app.py .
COPY qwen_classifier/ ./qwen_classifier/
COPY setup.py . 
RUN pip install --no-cache-dir -r requirements.txt

# Install PyTorch with CUDA support
RUN pip install --no-cache-dir \
    torch==2.1.2+cu121 \
    torchvision==0.16.2+cu121 \
    --extra-index-url https://download.pytorch.org/whl/cu121

# Set up Hugging Face authentication (use a build ARG for the token)
# ARG HF_TOKEN
# RUN python3 -c "from huggingface_hub import login; login(token='$HF_TOKEN')"

# Test model loading (use absolute import path)
RUN python3 -c "from qwen_classifier.model import QwenClassifier; \
    QwenClassifier.from_pretrained('KeivanR/Qwen2.5-1.5B-Instruct-MLB-clf_lora-1743189446'); \
    print('Model loaded successfully')"

# Run FastAPI app
EXPOSE 7860
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]