Duplicate from phamduyphuong251/test-deploy-dify

295ca33 2 days ago

2.71 kB

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)

	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)