Production-ready SNP model with tokenizer mapping and WSGI server

Dockerfile

@@ -1,27 +1,19 @@
 FROM python:3.10-slim
 
-# Make absolutely sure the working directory exists
-RUN mkdir -p /app
+# Create and switch to /app
 WORKDIR /app
 
-# Print working directory (for debugging)
-RUN echo "✅ Building from context:" && pwd && ls -R
-
-# Copy specific files explicitly by name
-COPY api_inference.py /app/api_inference.py
-COPY snp_universal_embedding.py /app/snp_universal_embedding.py
-COPY config.json /app/config.json
-COPY tokenizer.json /app/tokenizer.json
-COPY pytorch_model.bin /app/pytorch_model.bin
-COPY requirements.txt /app/requirements.txt
+# Copy all repo files into /app
+ADD . /app
 
 # Install dependencies
-RUN pip install --no-cache-dir -r /app/requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
 
 # Expose Hugging Face port
 EXPOSE 7860
 
-# Run the app
+# Environment variables for HF cache
 ENV HF_HOME=/tmp/huggingface
-ENV TRANSFORMERS_CACHE=/tmp/huggingface
-CMD ["python", "/app/api_inference.py"]
+
+# Use Gunicorn (production WSGI server)
+CMD ["gunicorn", "--bind", "0.0.0.0:7860", "api_inference:app"]

api_inference.py

@@ -14,7 +14,7 @@ from transformers import (
 # Environment Configuration
 # ============================================================
 os.environ["HF_HOME"] = "/tmp/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+# TRANSFORMERS_CACHE is deprecated; HF_HOME alone is enough
 
 MODEL_DIR = "./"
 PORT = int(os.environ.get("PORT", 7860))
@@ -53,6 +53,10 @@ class CustomSNPModel(PreTrainedModel):
 AutoConfig.register("custom_snp", CustomSNPConfig)
 AutoModel.register(CustomSNPConfig, CustomSNPModel)
 
+# --- Permanent Tokenizer Mapping ---
+from transformers import TOKENIZER_MAPPING, RobertaTokenizer
+TOKENIZER_MAPPING[CustomSNPConfig] = (RobertaTokenizer, RobertaTokenizer)
+
 
 # ============================================================
 # Load Model & Tokenizer
@@ -60,24 +64,28 @@ AutoModel.register(CustomSNPConfig, CustomSNPModel)
 try:
     print("Loading model from:", MODEL_DIR)
     config = AutoConfig.from_pretrained(MODEL_DIR, trust_remote_code=True)
-
-    # Try loading tokenizer; fallback if not mapped
-    from transformers import RobertaTokenizer
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
-    except Exception:
-        print("⚠️ Falling back to default RoBERTa tokenizer.")
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
-
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
     model = AutoModel.from_pretrained(MODEL_DIR, config=config, trust_remote_code=True)
     model.eval()
     print("✅ Custom SNP model loaded successfully.")
-
 except Exception as e:
     print("❌ Error loading custom model:", e)
     raise e
 
 
+# ============================================================
+# Initialize weights (optional for untrained layers)
+# ============================================================
+def initialize_weights_if_missing(model):
+    for name, param in model.named_parameters():
+        if param.requires_grad and (torch.isnan(param).any() or torch.all(param == 0)):
+            nn.init.xavier_uniform_(param)
+            print(f"🔧 Initialized missing weights: {name}")
+
+
+initialize_weights_if_missing(model)
+
+
 # ============================================================
 # Flask API Routes
 # ============================================================
@@ -122,7 +130,7 @@ def reason():
 
 
 # ============================================================
-# Run Server
+# Run Server (used by Gunicorn in production)
 # ============================================================
 if __name__ == "__main__":
     print(f"🚀 Starting SNP Universal Embedding API on port {PORT}")

requirements.txt

@@ -5,3 +5,4 @@ sentence-transformers
 flask
 numpy
 scikit-learn
+gunicorn