Create app.py

app.py

@@ -0,0 +1,186 @@
+"""
+HuggingFace Space - SPECTER2 Embedding API
+Academic paper embeddings using SPECTER2 with adapters
+"""
+
+import os
+from flask import Flask, request, jsonify
+from transformers import AutoTokenizer
+from adapters import AutoAdapterModel
+import torch
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Initialize Flask app
+app = Flask(__name__)
+
+# Load SPECTER2 model with adapters
+logger.info("Loading SPECTER2 base model and tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained('allenai/specter2_base')
+model = AutoAdapterModel.from_pretrained('allenai/specter2_base')
+
+logger.info("Loading SPECTER2 proximity adapter...")
+# Load the proximity adapter for similarity/retrieval tasks
+model.load_adapter("allenai/specter2", source="hf", load_as="specter2", set_active=True)
+logger.info("SPECTER2 model loaded successfully!")
+
+# Move to GPU if available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
+logger.info(f"Using device: {device}")
+
+
+def get_embeddings(texts):
+    """
+    Generate SPECTER2 embeddings for a list of texts
+
+    Args:
+        texts: List of strings (paper titles + abstracts)
+
+    Returns:
+        numpy array of embeddings (batch_size, 768)
+    """
+    # Tokenize
+    inputs = tokenizer(
+        texts,
+        padding=True,
+        truncation=True,
+        return_tensors="pt",
+        max_length=512
+    ).to(device)
+
+    # Generate embeddings
+    with torch.no_grad():
+        output = model(**inputs)
+        # Use [CLS] token embedding (first token)
+        embeddings = output.last_hidden_state[:, 0, :]
+
+    return embeddings.cpu().numpy()
+
+
+@app.route('/')
+def health():
+    """Health check endpoint"""
+    return jsonify({
+        "status": "healthy",
+        "model": "allenai/specter2",
+        "adapter": "proximity (similarity/retrieval)",
+        "dimensions": 768,
+        "device": str(device),
+        "endpoints": {
+            "/embed": "POST - Generate embedding for single text",
+            "/batch_embed": "POST - Generate embeddings for multiple texts"
+        }
+    })
+
+
+@app.route('/embed', methods=['POST'])
+def embed_text():
+    """
+    Generate embedding for a single text query
+
+    Request body:
+    {
+        "text": "Your paper title and abstract here"
+    }
+
+    Response:
+    {
+        "embedding": [0.123, -0.456, ...],
+        "dimensions": 768
+    }
+    """
+    try:
+        data = request.get_json()
+
+        if not data or 'text' not in data:
+            return jsonify({
+                "error": "Missing 'text' field in request body"
+            }), 400
+
+        text = data['text']
+
+        if not isinstance(text, str) or len(text.strip()) == 0:
+            return jsonify({
+                "error": "Text must be a non-empty string"
+            }), 400
+
+        # Generate embedding
+        embeddings = get_embeddings([text])
+
+        return jsonify({
+            "embedding": embeddings[0].tolist(),
+            "dimensions": len(embeddings[0])
+        })
+
+    except Exception as e:
+        logger.error(f"Error generating embedding: {str(e)}")
+        return jsonify({
+            "error": "Internal server error",
+            "message": str(e)
+        }), 500
+
+
+@app.route('/batch_embed', methods=['POST'])
+def batch_embed_texts():
+    """
+    Generate embeddings for multiple texts (batch processing)
+
+    Request body:
+    {
+        "texts": ["Paper 1 title and abstract", "Paper 2 title and abstract", ...]
+    }
+
+    Response:
+    {
+        "embeddings": [[0.123, ...], [0.456, ...], ...],
+        "count": 2,
+        "dimensions": 768
+    }
+    """
+    try:
+        data = request.get_json()
+
+        if not data or 'texts' not in data:
+            return jsonify({
+                "error": "Missing 'texts' field in request body"
+            }), 400
+
+        texts = data['texts']
+
+        if not isinstance(texts, list) or len(texts) == 0:
+            return jsonify({
+                "error": "Texts must be a non-empty list"
+            }), 400
+
+        # Limit batch size to prevent abuse
+        if len(texts) > 100:
+            return jsonify({
+                "error": "Batch size too large (max 100 texts)"
+            }), 400
+
+        # Generate embeddings
+        embeddings = get_embeddings(texts)
+
+        return jsonify({
+            "embeddings": embeddings.tolist(),
+            "count": len(embeddings),
+            "dimensions": embeddings.shape[1]
+        })
+
+    except Exception as e:
+        logger.error(f"Error generating batch embeddings: {str(e)}")
+        return jsonify({
+            "error": "Internal server error",
+            "message": str(e)
+        }), 500
+
+
+if __name__ == '__main__':
+    # HuggingFace Spaces requires the app to listen on port 7860
+    port = int(os.environ.get('PORT', 7860))
+    logger.info(f"Starting server on port {port}...")
+    app.run(host='0.0.0.0', port=port, debug=False)