wehub-resource-sync
55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
import torch
|
|
from transformers import BertModel, BertTokenizerFast, pipeline
|
|
|
|
import mlflow
|
|
|
|
sentence_transformers_architecture = "sentence-transformers/all-MiniLM-L12-v2"
|
|
task = "feature-extraction"
|
|
|
|
model = BertModel.from_pretrained(sentence_transformers_architecture)
|
|
tokenizer = BertTokenizerFast.from_pretrained(sentence_transformers_architecture)
|
|
|
|
sentence_transformer_pipeline = pipeline(task=task, model=model, tokenizer=tokenizer)
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.transformers.log_model(
|
|
transformers_model=sentence_transformer_pipeline,
|
|
name="sentence_transformer",
|
|
framework="pt",
|
|
torch_dtype=torch.bfloat16,
|
|
)
|
|
|
|
loaded = mlflow.transformers.load_model(model_info.model_uri, return_type="components")
|
|
|
|
|
|
def pool_and_normalize_encodings(input_sentences, model, tokenizer, **kwargs):
|
|
def pool(model_output, attention_mask):
|
|
embeddings = model_output[0]
|
|
expanded_mask = attention_mask.unsqueeze(-1).expand(embeddings.size()).float()
|
|
return torch.sum(embeddings * expanded_mask, 1) / torch.clamp(
|
|
expanded_mask.sum(1), min=1e-9
|
|
)
|
|
|
|
encoded = tokenizer(
|
|
input_sentences,
|
|
padding=True,
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
)
|
|
with torch.no_grad():
|
|
model_output = model(**encoded)
|
|
|
|
pooled = pool(model_output, encoded["attention_mask"])
|
|
return torch.nn.functional.normalize(pooled, p=2, dim=1)
|
|
|
|
|
|
sentences = [
|
|
"He said that he's sinking all of his investment budget into coconuts.",
|
|
"No matter how deep you dig, there's going to be a point when it just gets too hot.",
|
|
"She said that there isn't a noticeable difference between a 10 year and a 15 year whisky.",
|
|
]
|
|
|
|
encoded_sentences = pool_and_normalize_encodings(sentences, **loaded)
|
|
|
|
print(encoded_sentences)
|