File size: 1,231 Bytes
5901795 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
from transformers import AutoModel, AutoTokenizer
import torch
class EndpointHandler():
def __init__(self, path=""):
# Initialize the tokenizer and model with pre-trained weights
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModel.from_pretrained(path)
def __call__(self, data):
# Extract text input from the request data
inputs = data['inputs']
# Define a prompt to provide context
prompt = "Contextual understanding of the following text, from the perspective of Chassidic philosophy: "
# Combine prompt with the actual input
combined_input = prompt + inputs
# Prepare the text for the model
encoded_input = self.tokenizer(combined_input, return_tensors='pt', padding=True, truncation=True, max_length=512)
# Generate embeddings without updating gradients
with torch.no_grad():
outputs = self.model(**encoded_input)
# Extract embeddings from the last hidden layer
embeddings = outputs.last_hidden_state.squeeze().tolist()
# Return the embeddings as a list (serialized format)
return {'embeddings': embeddings}
|