Create model.py

model.py

@@ -0,0 +1,21 @@
+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 chunks from the payload
+        chunks = payload.get("chunks", [])
+        
+        # Generate embeddings for chunks
+        embeddings = self.embedding_model.encode(chunks, convert_to_tensor=True)
+        
+        # Prepare response with tensor embeddings
+        response = {
+            "embeddings": embeddings.tolist(),  # Convert tensor to list for JSON serialization
+            "shape": list(embeddings.shape)    # Return the shape of the embeddings tensor
+        }
+        return response