Update api_inference.py

api_inference.py

@@ -6,7 +6,7 @@ from transformers import (
     AutoTokenizer,
     AutoModel,
     AutoConfig,
-    PretrainedConfig,
+    PretrainedConfig,  # <-- 1. YOURS: Using PretrainedConfig (This is correct)
     PreTrainedModel,
 )
 
@@ -25,18 +25,23 @@ PORT = int(os.environ.get("PORT", 7860))
 app = Flask(__name__)
 
 # ============================================================
-# Register Custom SNP Architecture (THE FIX IS HERE)
+# Register Custom SNP Architecture
 # ============================================================
-class CustomSNPConfig(AutoConfig):
-    # This will correctly inherit 'custom_snp' from your config.json
+
+# This is from YOUR original file. It is correct because
+# PretrainedConfig has the .from_dict() method.
+class CustomSNPConfig(PretrainedConfig):
     model_type = "custom_snp"
 
+
+# This is from MY file. This is the fix for your 500 error,
+# as it loads the real transformer model.
 class CustomSNPModel(PreTrainedModel):
     config_class = CustomSNPConfig
 
     def __init__(self, config):
         super().__init__(config)
-        # This is the correct way to load the base transformer
+        # This correctly loads the base transformer
         self.shared_encoder = AutoModel.from_config(config)
         
         hidden_size = self.shared_encoder.config.hidden_size
@@ -63,6 +68,7 @@ class CustomSNPModel(PreTrainedModel):
         proj = self.projection(x)
         return proj # Return the final projection
 
+
 # Register model so AutoModel recognizes it
 AutoConfig.register("custom_snp", CustomSNPConfig)
 AutoModel.register(CustomSNPConfig, CustomSNPModel)
@@ -75,8 +81,12 @@ try:
     print("Loading model from:", MODEL_DIR)
     
     # trust_remote_code=True is essential for this to work
+    # This will now succeed because CustomSNPConfig has .from_dict()
     config = AutoConfig.from_pretrained(MODEL_DIR, trust_remote_code=True)
+    
     tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
+    
+    # This will load your custom model architecture
     model = AutoModel.from_pretrained(MODEL_DIR, config=config, trust_remote_code=True)
     
     model.eval()
@@ -89,7 +99,7 @@ except Exception as e:
 
 
 # ============================================================
-# Flask API Routes (Your routes are correct)
+# Flask API Routes
 # ============================================================
 @app.route("/", methods=["GET"])
 def home():
@@ -105,19 +115,19 @@ def health():
 def embed():
     try:
         data = request.get_json(force=True)
-        text = data.get("text", "")
+        text = data.get("text", "") # Your Colab script is correct, it sends "text"
         if not text:
             return jsonify({"error": "Text is required"}), 400
 
         # Tokenize the text
         inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
         
-        # Run inference
+        # Run inference (no .float() conversion needed)
         with torch.no_grad():
             embeddings = model(**inputs)
         
-        # The model's forward() method now directly returns the projection
-        return jsonify({"embedding": embeddings.tolist()})
+        # Return the projection
+        return jsonify({"embedding": embeddings.tolist()}) # Your Colab script is correct, it expects "embedding"
 
     except Exception as e:
         print(f"ERROR in /embed: {e}")