KernelMemory.Extensions/src/python/sbert.py at develop · alkampfergit/KernelMemory.Extensions · GitHub

History

93 lines (74 loc) · 2.99 KB

Raw

from fastapi import FastAPI

import uvicorn

import sys

import argparse

import os

from sentence_transformers import SentenceTransformer

from fastapi import HTTPException

from pydantic import BaseModel

from typing import List

import sentence_transformers_pool

import logging

# Configure logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

transformers_cache = os.environ.get('TRANSFORMERS_CACHE')

# Pool dictionary will create a pool of sentence transformers for each model

# we need to manage not too many models at the same time due to memory constraints

# for this example we will not limit the number of models

transformers_dict = {}

app = FastAPI()

class SentenceInput(BaseModel):

sentences: List[str]

modelName: str = "all-mpnet-base-v2"

class DimensionsInput(BaseModel):

modelName: str = "all-mpnet-base-v2"

class CountTokenInput(BaseModel):

modelName: str = "all-mpnet-base-v2"

text: str

def load_model(model_name: str) -> SentenceTransformer:

if model_name not in transformers_dict:

transformers_dict[model_name] = sentence_transformers_pool.SentenceTransformerPool(model_name=model_name, max_size=3)

return transformers_dict[model_name].get()

def release_model(model_name: str, model: SentenceTransformer) -> None:

transformers_dict[model_name].release(model)

@app.get("/ping")

async def ping():

return "ping"

@app.post("/dimensions")

async def dimensions(input_data: DimensionsInput):

model = load_model(input_data.modelName)

try:

return { "model" : input_data.modelName,

"maxSequenceLength" : model.max_seq_length,

"dimension" : model.get_sentence_embedding_dimension() }

finally:

release_model(input_data.modelName, model)

@app.post("/process-sentences")

async def process_sentences(input_data: SentenceInput):

sentences = input_data.sentences

model = load_model(input_data.modelName)

try:

encoded = model.encode(sentences)

return {"embeddings": encoded.tolist(), "model" : input_data.modelName}

finally:

release_model(input_data.modelName, model)

@app.post("/count-tokens")

async def count_tokens(input_data: CountTokenInput):

model = load_model(input_data.modelName)

try:

tokenizer = model.tokenizer

tokens = tokenizer.tokenize(input_data.text)

return { "count" : len(tokens) }

finally:

logging.info("Releasing model")

release_model(input_data.modelName, model)

if __name__ == "__main__":

# Create an argument parser

parser = argparse.ArgumentParser()

parser.add_argument("--port", type=int, default=8000, help="Specify the port number")

# Parse the command-line arguments

args = parser.parse_args()

# Get the port number from the parsed arguments

port = args.port

if __name__ == "__main__":

uvicorn.run(app, host="0.0.0.0", port=port)

KernelMemory.Extensions/src/python/sbert.py at develop · alkampfergit/KernelMemory.Extensions · GitHub

Expand file tree

Breadcrumbs

sbert.py

Latest commit

History

Breadcrumbs

sbert.py

File metadata and controls