366degrees/snp-universal-embedding

1import os
2import torch
3import torch.nn as nn
4from flask import Flask, request, jsonify
5from transformers import (
6    AutoTokenizer,
7    AutoModel,
8    AutoConfig,
9    PretrainedConfig,
10    PreTrainedModel,
11)
12
13# ============================================================
14# Redirect Hugging Face cache to /app/hf_cache (always writable)
15CACHE_DIR = "/app/hf_cache"
16os.makedirs(CACHE_DIR, exist_ok=True)
17os.environ["HF_HOME"] = CACHE_DIR
18os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
19os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
20
21MODEL_DIR = "./"
22PORT = int(os.environ.get("PORT", 7860))
23
24app = Flask(__name__)
25
26
27# ============================================================
28# Register Custom SNP Architecture
29# ============================================================
30class CustomSNPConfig(PretrainedConfig):
31    model_type = "custom_snp"
32
33
34class CustomSNPModel(PreTrainedModel):
35    config_class = CustomSNPConfig
36
37    def __init__(self, config):
38        super().__init__(config)
39        hidden_size = getattr(config, "hidden_size", 768)
40        # Mirror and Prism heads
41        self.encoder = nn.Linear(hidden_size, hidden_size)
42        self.mirror_head = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.Tanh())
43        self.prism_head = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.Tanh())
44        self.projection = nn.Linear(hidden_size, 6)
45
46    def forward(self, input_ids=None, attention_mask=None, **kwargs):
47        # Simulate encoded representations
48        x = self.encoder(input_ids.float()) if input_ids is not None else None
49        x = self.mirror_head(x)
50        x = self.prism_head(x)
51        return self.projection(x)
52
53
54# Register model so AutoModel recognizes it
55AutoConfig.register("custom_snp", CustomSNPConfig)
56AutoModel.register(CustomSNPConfig, CustomSNPModel)
57
58
59# ============================================================
60# Load Model & Tokenizer
61# ============================================================
62try:
63    print("Loading model from:", MODEL_DIR)
64    config = AutoConfig.from_pretrained(MODEL_DIR, trust_remote_code=True)
65
66    # Try loading tokenizer; fallback if not mapped
67    from transformers import RobertaTokenizer
68    try:
69        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
70    except Exception:
71        print("⚠️ Falling back to default RoBERTa tokenizer.")
72        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
73
74    model = AutoModel.from_pretrained(MODEL_DIR, config=config, trust_remote_code=True)
75    model.eval()
76    print("✅ Custom SNP model loaded successfully.")
77
78except Exception as e:
79    print("❌ Error loading custom model:", e)
80    raise e
81
82
83# ============================================================
84# Flask API Routes
85# ============================================================
86@app.route("/", methods=["GET"])
87def home():
88    return jsonify({"status": "SNP Universal Embedding API running"})
89
90
91@app.route("/health", methods=["GET"])
92def health():
93    return jsonify({"status": "healthy"})
94
95
96@app.route("/embed", methods=["POST"])
97def embed():
98    data = request.get_json(force=True)
99    text = data.get("text", "")
100    if not text:
101        return jsonify({"error": "Text is required"}), 400
102
103    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
104    with torch.no_grad():
105        embeddings = model(**inputs)
106    if hasattr(embeddings, "last_hidden_state"):
107        embeddings = embeddings.last_hidden_state.mean(dim=1)
108    elif isinstance(embeddings, tuple):
109        embeddings = embeddings[0]
110    return jsonify({"embedding": embeddings.tolist()})
111
112
113@app.route("/reason", methods=["POST"])
114def reason():
115    data = request.get_json(force=True)
116    premise = data.get("premise", "")
117    hypothesis = data.get("hypothesis", "")
118    combined = f"{premise} {hypothesis}"
119    inputs = tokenizer(combined, return_tensors="pt", truncation=True, padding=True)
120    with torch.no_grad():
121        output = model(**inputs)
122    score = float(output.mean().item())
123    return jsonify({"reasoning_score": score})
124
125
126# ============================================================
127# Run Server
128# ============================================================
129if __name__ == "__main__":
130    print(f"🚀 Starting SNP Universal Embedding API on port {PORT}")
131    app.run(host="0.0.0.0", port=PORT)
132
133