File size: 2,705 Bytes
295ca33 | 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | from typing import Dict, List, Union
import torch
from transformers import AutoModel
from custom_tokenizer import CustomPhobertTokenizer
def mean_pooling(model_output, attention_mask):
token_embeddings = model_output[
0
] # First element of model_output contains all token embeddings
input_mask_expanded = (
attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
)
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
input_mask_expanded.sum(1), min=1e-9
)
class PreTrainedPipeline:
def __init__(self, path="."):
self.model = AutoModel.from_pretrained(path)
self.tokenizer = CustomPhobertTokenizer.from_pretrained(path)
def __call__(self, inputs: Dict[str, Union[str, List[str]]]) -> List[float]:
"""
Args:
inputs (Dict[str, Union[str, List[str]]]):
a dictionary containing a query sentence and a list of key sentences
"""
# Combine the query sentence and key sentences into one list
sentences = [inputs["source_sentence"]] + inputs["sentences"]
# Tokenize sentences
encoded_input = self.tokenizer(
sentences, padding=True, truncation=True, return_tensors="pt"
)
# Compute token embeddings
with torch.no_grad():
model_output = self.model(**encoded_input)
# Perform pooling to get sentence embeddings
sentence_embeddings = mean_pooling(
model_output, encoded_input["attention_mask"]
)
# Separate the query embedding from the key embeddings
query_embedding = sentence_embeddings[0]
key_embeddings = sentence_embeddings[1:]
# Compute cosine similarities (or any other comparison method you prefer)
cosine_similarities = torch.nn.functional.cosine_similarity(
query_embedding.unsqueeze(0), key_embeddings
)
# Convert the tensor of cosine similarities to a list of floats
scores = cosine_similarities.tolist()
return scores
if __name__ == "__main__":
inputs = {
"source_sentence": "Anh ấy đang là sinh viên năm cuối",
"sentences": [
"Anh ấy học tại Đại học Bách khoa Hà Nội, chuyên ngành Khoa học máy tính",
"Anh ấy đang làm việc tại nhà máy sản xuất linh kiện điện tử",
"Anh ấy chuẩn bị đi du học nước ngoài",
"Anh ấy sắp mở cửa hàng bán mỹ phẩm",
"Nhà anh ấy có rất nhiều cây cảnh",
],
}
pipeline = PreTrainedPipeline()
res = pipeline(inputs)
|