Update app.py

app.py

@@ -1,11 +1,13 @@
 """
-HuggingFace Space - Embedding API
-Lightweight stateless API for generating text embeddings
+HuggingFace Space - SPECTER2 Embedding API
+Academic paper embeddings using SPECTER2 with adapters
 """
 
 import os
 from flask import Flask, request, jsonify
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer
+from adapters import AutoAdapterModel
+import torch
 import logging
 
 # Configure logging
@@ -15,13 +17,49 @@ logger = logging.getLogger(__name__)
 # Initialize Flask app
 app = Flask(__name__)
 
-# Load the embedding model (SPECTER2, 768 dimensions)
-# SPECTER2 is trained on scientific papers with citation relationships
-# Best-in-class for academic paper embeddings
-logger.info("Loading SPECTER2 embedding model...")
-model = SentenceTransformer('allenai/specter2')
+# 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():
@@ -29,9 +67,11 @@ def health():
     return jsonify({
         "status": "healthy",
         "model": "allenai/specter2",
+        "adapter": "proximity (similarity/retrieval)",
         "dimensions": 768,
+        "device": str(device),
         "endpoints": {
-            "/embed": "POST - Generate embeddings for text",
+            "/embed": "POST - Generate embedding for single text",
             "/batch_embed": "POST - Generate embeddings for multiple texts"
         }
     })
@@ -44,7 +84,7 @@ def embed_text():
 
     Request body:
     {
-        "text": "Your text here"
+        "text": "Your paper title and abstract here"
     }
 
     Response:
@@ -69,11 +109,11 @@ def embed_text():
             }), 400
 
         # Generate embedding
-        embedding = model.encode(text, convert_to_numpy=True)
+        embeddings = get_embeddings([text])
 
         return jsonify({
-            "embedding": embedding.tolist(),
-            "dimensions": len(embedding)
+            "embedding": embeddings[0].tolist(),
+            "dimensions": len(embeddings[0])
         })
 
     except Exception as e:
@@ -91,7 +131,7 @@ def batch_embed_texts():
 
     Request body:
     {
-        "texts": ["Text 1", "Text 2", ...]
+        "texts": ["Paper 1 title and abstract", "Paper 2 title and abstract", ...]
     }
 
     Response:
@@ -123,12 +163,12 @@ def batch_embed_texts():
             }), 400
 
         # Generate embeddings
-        embeddings = model.encode(texts, convert_to_numpy=True)
+        embeddings = get_embeddings(texts)
 
         return jsonify({
             "embeddings": embeddings.tolist(),
             "count": len(embeddings),
-            "dimensions": embeddings.shape[1] if len(embeddings) > 0 else 768
+            "dimensions": embeddings.shape[1]
         })
 
     except Exception as e: