Spaces:
Build error
Build error
File size: 2,110 Bytes
4589b64 |
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 |
from src.model.model import CommentGenerator
from src.model.text_process import TextPreprocess
class GetComment():
def __init__(self) -> None:
"""Get Comment init funtion
"""
self.CommentGenerator = CommentGenerator()
self.text_preprocess = TextPreprocess()
self.inputs = {}
def generator(self, info):
"""Generation function
Args:
info (Dict): info extract from kafka message
Returns:
List: List of comment
"""
content = info["content"]
content = self.text_preprocess.preprocess(content)
medias = info["medias"]
if not medias:
medias = [None]
print(medias)
typfeed = info["type_generation"]
num_comments = info["num_comments"]
mapper = {
0: "bảng tin",
1: "trải nghiệm"
}
comments_prefix = ["thứ một",
"thứ hai",
"thứ ba"]
comments = []
for i, comment_prefix in enumerate(comments_prefix):
content_w_prefix = f"{mapper[typfeed]}: {comment_prefix}: {content}"
self.inputs[f"content_{i}"] = self.CommentGenerator.get_text_feature(content_w_prefix)
while len(comments) < num_comments:
for i, media in enumerate(medias):
print(i)
if i not in self.inputs:
self.inputs[i] = self.CommentGenerator.get_image_feature_from_url(media, is_local=True)
image_feature, image_mask = self.inputs[i]
for i in range(len(comments_prefix)):
content_feature, content_mask = self.inputs[f"content_{i}"]
comment = self.CommentGenerator.inference(content_feature, content_mask, image_feature, image_mask)[0]
comments.append(comment)
comments = list(set(comments))
if len(comments) >= num_comments:
return comments
|