Simplified: Flask direct run, no Gunicorn

Dockerfile

@@ -27,4 +27,5 @@ RUN touch /.gitconfig && chmod 666 /.gitconfig
 USER appuser
 
 # 8) Start server (Gunicorn)
-CMD ["gunicorn", "--bind", "0.0.0.0:7860", "api_inference:app"]
+CMD ["python", "api_inference.py"]
+

api_inference.py

@@ -3,19 +3,16 @@ import torch
 import torch.nn as nn
 from flask import Flask, request, jsonify
 from transformers import (
-    AutoTokenizer,
-    AutoModel,
     AutoConfig,
+    AutoModel,
     PretrainedConfig,
     PreTrainedModel,
-    TOKENIZER_MAPPING,
-    RobertaTokenizer
 )
+from transformers import RobertaTokenizerFast as RobertaTokenizer
 
 # ============================================================
-# Early Custom Registration (Runs before Flask/Gunicorn)
+# Custom SNP Architecture (no Gunicorn complications)
 # ============================================================
-
 class CustomSNPConfig(PretrainedConfig):
     model_type = "custom_snp"
 
@@ -40,47 +37,42 @@ class CustomSNPModel(PreTrainedModel):
         return self.projection(x)
 
 
-# --- Force registration order ---
-TOKENIZER_MAPPING[CustomSNPConfig] = (RobertaTokenizer, RobertaTokenizer)
-AutoConfig.register("custom_snp", CustomSNPConfig)
-AutoModel.register(CustomSNPConfig, CustomSNPModel)
-
 # ============================================================
-# Environment Configuration
+# Environment
 # ============================================================
 os.environ["HF_HOME"] = "/tmp/huggingface"
 MODEL_DIR = "./"
 PORT = int(os.environ.get("PORT", 7860))
-
 app = Flask(__name__)
 
 # ============================================================
-# Load Model & Tokenizer
+# Load Model & Tokenizer (direct tokenizer, no AutoTokenizer)
 # ============================================================
 try:
     print("Loading model from:", MODEL_DIR)
     config = AutoConfig.from_pretrained(MODEL_DIR, trust_remote_code=True)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
-    model = AutoModel.from_pretrained(MODEL_DIR, config=config, trust_remote_code=True)
+
+    # Use concrete tokenizer to avoid mapping issues
+    try:
+        tokenizer = RobertaTokenizer.from_pretrained(MODEL_DIR)
+        print("✅ Loaded tokenizer from model directory.")
+    except Exception:
+        print("⚠️ Falling back to default roberta-base tokenizer.")
+        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+
+    model = CustomSNPModel(config)
+    if os.path.exists(os.path.join(MODEL_DIR, "pytorch_model.bin")):
+        state = torch.load(os.path.join(MODEL_DIR, "pytorch_model.bin"), map_location="cpu")
+        model.load_state_dict(state, strict=False)
     model.eval()
     print("✅ Custom SNP model loaded successfully.")
 except Exception as e:
     print("❌ Error loading custom model:", e)
     raise e
 
-# ============================================================
-# Initialize weights (optional)
-# ============================================================
-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 Routes
+# Routes
 # ============================================================
 @app.route("/", methods=["GET"])
 def home():
@@ -135,8 +127,9 @@ def test():
         "embedding_preview": vector[0][:6]
     })
 
+
 # ============================================================
-# Run Server (Gunicorn)
+# Run Flask directly (no Gunicorn)
 # ============================================================
 if __name__ == "__main__":
     print(f"🚀 Starting SNP Universal Embedding API on port {PORT}")