Update ML service to use transformers instead of sentence-transformers for compatibility

app.py

@@ -1,11 +1,25 @@
 from flask import Flask, request, jsonify
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
+import torch
 import os
 
 app = Flask(__name__)
 
 # Load your model once
-model = SentenceTransformer("all-MiniLM-L6-v2")
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
+
+def get_embedding(text):
+    """Generate embedding for a single text"""
+    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        # Use mean pooling over token embeddings
+        embeddings = outputs.last_hidden_state.mean(dim=1)
+        # Normalize the embeddings
+        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+    return embeddings.squeeze().tolist()
 
 @app.route('/api/predict', methods=['POST'])
 def predict():
@@ -18,13 +32,16 @@ def predict():
         if not isinstance(texts, list):
             return jsonify({'error': 'Data must be a list of texts'}), 400
 
-        # Generate embeddings
-        embeddings = model.encode(texts, normalize_embeddings=True)
+        # Generate embeddings for each text
+        embeddings = []
+        for text in texts:
+            if isinstance(text, str):
+                embedding = get_embedding(text)
+                embeddings.append(embedding)
+            else:
+                return jsonify({'error': 'All items in data must be strings'}), 400
 
-        # Convert to list format
-        embeddings_list = embeddings.tolist() if hasattr(embeddings, 'tolist') else embeddings
-
-        return jsonify({'data': embeddings_list})
+        return jsonify({'data': embeddings})
 
     except Exception as e:
         return jsonify({'error': str(e)}), 500

requirements.txt

@@ -1,4 +1,4 @@
 Flask==2.3.3
-sentence-transformers==2.2.2
+transformers==4.36.0
 torch>=2.0.0
 numpy>=1.21.0