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 text chunks from the payload
        chunks = payload.get("inputs", [])

        # 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