Create model.py

model.py

@@ -0,0 +1,25 @@
+from sentence_transformers import SentenceTransformer
+import torch
+
+class Model:
+    def __init__(self):
+        # Load the pre-trained model
+        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+    def __call__(self, payload):
+        # Extract inputs from the payload
+        inputs = payload.get("inputs", {})
+        source_sentence = inputs.get("source_sentence", "")
+        sentences = inputs.get("sentences", [])
+
+        # Combine source_sentence with sentences
+        chunks = [source_sentence] + sentences
+        # Generate embeddings
+        embeddings = self.embedding_model.encode(chunks, convert_to_tensor=True)
+
+        # Prepare response
+        response = {
+            "embeddings": embeddings.tolist(),  # Convert tensor to list for JSON serialization
+            "shape": list(embeddings.shape)    # Return the shape of the embeddings tensor
+        }
+        return response