Fix tokenizer mapping for CustomSNPConfig

api_inference.py

@@ -3,15 +3,27 @@ import torch
 import torch.nn as nn
 from flask import Flask, request, jsonify
 from transformers import (
-    AutoConfig,
+    AutoTokenizer,
     AutoModel,
+    AutoConfig,
     PretrainedConfig,
     PreTrainedModel,
 )
-from transformers import RobertaTokenizerFast as RobertaTokenizer
 
 # ============================================================
-# Custom SNP Architecture (no Gunicorn complications)
+# Environment Configuration
+# ============================================================
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+
+MODEL_DIR = "./"
+PORT = int(os.environ.get("PORT", 7860))
+
+app = Flask(__name__)
+
+
+# ============================================================
+# Register Custom SNP Architecture
 # ============================================================
 class CustomSNPConfig(PretrainedConfig):
     model_type = "custom_snp"
@@ -22,72 +34,70 @@ class CustomSNPModel(PreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
-        hidden = getattr(config, "hidden_size", 768)
-        self.encoder = nn.Linear(hidden, hidden)
-        self.mirror_head = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh())
-        self.prism_head = nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh())
-        self.projection = nn.Linear(hidden, 6)
+        hidden_size = getattr(config, "hidden_size", 768)
+        # Mirror and Prism heads
+        self.encoder = nn.Linear(hidden_size, hidden_size)
+        self.mirror_head = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.Tanh())
+        self.prism_head = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.Tanh())
+        self.projection = nn.Linear(hidden_size, 6)
 
     def forward(self, input_ids=None, attention_mask=None, **kwargs):
-        if input_ids is None:
-            raise ValueError("input_ids required")
-        x = self.encoder(input_ids.float())
+        # Simulate encoded representations
+        x = self.encoder(input_ids.float()) if input_ids is not None else None
         x = self.mirror_head(x)
         x = self.prism_head(x)
         return self.projection(x)
 
 
-# ============================================================
-# Environment
-# ============================================================
-os.environ["HF_HOME"] = "/tmp/huggingface"
-MODEL_DIR = "./"
-PORT = int(os.environ.get("PORT", 7860))
-app = Flask(__name__)
+# Register model so AutoModel recognizes it
+AutoConfig.register("custom_snp", CustomSNPConfig)
+AutoModel.register(CustomSNPConfig, CustomSNPModel)
+
 
 # ============================================================
-# Load Model & Tokenizer (direct tokenizer, no AutoTokenizer)
+# Load Model & Tokenizer
 # ============================================================
 try:
     print("Loading model from:", MODEL_DIR)
     config = AutoConfig.from_pretrained(MODEL_DIR, trust_remote_code=True)
 
-    # Use concrete tokenizer to avoid mapping issues
+    # Try loading tokenizer; fallback if not mapped
+    from transformers import RobertaTokenizer
     try:
-        tokenizer = RobertaTokenizer.from_pretrained(MODEL_DIR)
-        print("✅ Loaded tokenizer from model directory.")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
     except Exception:
-        print("⚠️ Falling back to default roberta-base tokenizer.")
+        print("⚠️ Falling back to default RoBERTa 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 = 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
 
 
 # ============================================================
-# Routes
+# Flask API Routes
 # ============================================================
 @app.route("/", methods=["GET"])
 def home():
     return jsonify({"status": "SNP Universal Embedding API running"})
 
+
 @app.route("/health", methods=["GET"])
 def health():
     return jsonify({"status": "healthy"})
 
+
 @app.route("/embed", methods=["POST"])
 def embed():
     data = request.get_json(force=True)
     text = data.get("text", "")
     if not text:
         return jsonify({"error": "Text is required"}), 400
+
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         embeddings = model(**inputs)
@@ -97,6 +107,7 @@ def embed():
         embeddings = embeddings[0]
     return jsonify({"embedding": embeddings.tolist()})
 
+
 @app.route("/reason", methods=["POST"])
 def reason():
     data = request.get_json(force=True)
@@ -109,28 +120,11 @@ def reason():
     score = float(output.mean().item())
     return jsonify({"reasoning_score": score})
 
-@app.route("/test", methods=["GET"])
-def test():
-    sample_text = "She knows he cheats but stays anyway."
-    inputs = tokenizer(sample_text, return_tensors="pt")
-    with torch.no_grad():
-        output = model(**inputs)
-    if hasattr(output, "last_hidden_state"):
-        vector = output.last_hidden_state.mean(dim=1).tolist()
-    elif isinstance(output, tuple):
-        vector = output[0].tolist()
-    else:
-        vector = output.tolist()
-    return jsonify({
-        "message": "SNP Universal Embedding model is active.",
-        "sample_text": sample_text,
-        "embedding_preview": vector[0][:6]
-    })
-
 
 # ============================================================
-# Run Flask directly (no Gunicorn)
+# Run Server
 # ============================================================
 if __name__ == "__main__":
     print(f"🚀 Starting SNP Universal Embedding API on port {PORT}")
     app.run(host="0.0.0.0", port=PORT)
+