init

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+*test*

app.py

@@ -0,0 +1,60 @@
+import streamlit as st
+from PIL import Image
+
+#Trick to not init function multitime
+if "get_comment" not in st.session_state:
+    print("INIT MODEL")
+    from src.comment import GetComment
+    st.session_state.get_comment = GetComment()
+    print("DONE INIT MODEL")
+
+st.set_page_config(page_title="Auto Comment Generation", layout="wide", page_icon = "./storage/linhai.jpeg")
+hide_menu_style = """
+<style>
+footer {visibility: hidden;}
+</style>
+"""
+st.markdown(hide_menu_style, unsafe_allow_html= True)
+
+st.markdown(
+    """
+    <style>
+    [data-testid="stSidebar"][aria-expanded="true"] > div:first-child{
+        width: 400px;
+    }
+    [data-testid="stSidebar"][aria-expanded="false"] > div:first-child{
+        margin-left: -400px;
+    }
+     
+    """,
+    unsafe_allow_html=True,
+)
+
+left_col, right_col = st.columns(2)
+
+#Sidebar
+num_comments = st.sidebar.slider("Number of comments generated", 1, 10, 1)
+type_generation = st.sidebar.radio("Type of comment generation", options=[0,1])
+
+#LEFT COLUMN
+with left_col:
+    content = st.text_input("Enter your status", value="Hôm nay trời đẹp quá nhỉ")
+    images = st.file_uploader("Choose an image file", key=1, accept_multiple_files=True, type=["jpg", "jpeg", "png", "webp", ])
+    
+if left_col.button("Generate"):
+    info = {
+        "content": content,
+        "medias": images,
+        "type_generation": type_generation,
+        "num_comments": num_comments
+    }
+    
+    #RIGHT COLUMN
+    with right_col:
+        st.write("**STATUS:** ", info["content"])
+        st.write("**IMAGES:** ")
+        st.image(info["medias"], width=100, caption=list(range(1, len(images) +1)))
+        
+        comments = st.session_state.get_comment.generator(info)
+        for i, comment in enumerate(comments):
+            st.write(f"**Comment {i+1}**: {comment}")        

config/comment_generator.yaml

@@ -0,0 +1,7 @@
+model:
+  device: cpu
+  url: https://drive.google.com/uc?id=1-Fyh5nBbbUKwVayg3USuSYtjJQcmuFrV
+  dir: ./storage/last_vit_base_vit5_base_input_256_output_32_v4_thresh_2.0.pt
+  input_maxlen: 256
+  output_maxlen: 32
+  num_comment: 10

requirements.txt

@@ -0,0 +1,5 @@
+torch==2.0.1
+gdown
+transformers==4.32.0
+Pillow==9.5.0
+streamlit==1.26.0

src/comment.py

@@ -0,0 +1,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
+        

src/model/init.py

@@ -0,0 +1,7 @@
+import os
+import gdown
+
+def download_model(url, output):
+    if os.path.isfile(output):
+        return
+    gdown.download(url, output, quiet=True)

src/model/model.py

@@ -0,0 +1,79 @@
+import os
+import torch
+import requests
+from PIL import Image
+
+from utils.config import Config
+from src.model.init import download_model
+from transformers import AutoImageProcessor, ViTModel, AutoTokenizer, T5EncoderModel
+
+class CommentGenerator():
+    def __init__(self) -> None:
+        
+        self.config = Config("./config/comment_generator.yaml").__get_config__()
+        download_model(self.config['model']['url'], self.config['model']['dir'])
+        
+        #Get model
+        self.tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base")
+        self.model = torch.load(self.config["model"]["dir"], map_location=torch.device(self.config["model"]['device']))        
+        
+        
+        #Image
+        self.vit_image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+        self.vit_model = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k")
+        self.vit_model.to(self.config["model"]["device"])
+        
+        #Content
+        self.vit5_model = T5EncoderModel.from_pretrained("VietAI/vit5-base")
+        self.vit5_model.to(self.config["model"]["device"])
+        
+    def get_text_feature(self, content):
+        inputs = self.tokenizer(content,
+                            padding="max_length",
+                            truncation=True,
+                            max_length=self.config["model"]["input_maxlen"],
+                            return_tensors="pt").to(self.config["model"]["device"])
+        with torch.no_grad():
+            outputs = self.vit5_model(**inputs)
+        last_hidden_states = outputs.last_hidden_state
+        return last_hidden_states.to(self.config["model"]["device"]), inputs.attention_mask.to(self.config["model"]["device"])
+        
+
+    def get_image_feature_from_url(self, image_url, is_local=False):
+        if not image_url:
+            print(f"WARNING not image url {image_url}")
+            return torch.zeros((1, 197, 768)).to(self.config["model"]["device"]), torch.zeros((1, 197)).to(self.config["model"]["device"])
+        if not is_local:
+            try:
+                images = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
+            except:
+                print(f"READ IMAGE ERR: {image_url}")
+                return torch.zeros((1, 197, 768)).to(self.config["model"]["device"]), torch.zeros((1, 197)).to(self.config["model"]["device"])
+        else:
+            images = Image.open(image_url).convert("RGB")
+        inputs = self.vit_image_processor(images, return_tensors="pt").to(self.config["model"]["device"])
+        with torch.no_grad():
+            outputs = self.vit_model(**inputs)
+        last_hidden_states = outputs.last_hidden_state
+        attention_mask = torch.ones((last_hidden_states.shape[0], last_hidden_states.shape[1]))
+
+        return last_hidden_states.to(self.config["model"]["device"]), attention_mask.to(self.config["model"]["device"])
+
+    def inference(self, content_feature, content_mask, image_feature, image_mask):
+        
+        inputs_embeds = torch.cat((image_feature[0], content_feature[0]), 0)
+        inputs_embeds = torch.unsqueeze(inputs_embeds, 0)
+        attention_mask = torch.cat((image_mask[0], content_mask[0]), 0)
+        attention_mask = torch.unsqueeze(attention_mask, 0)
+        with torch.no_grad():
+            generated_ids = self.model.generate(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                num_beams=2,
+                max_length=self.config["model"]["output_maxlen"],
+                # num_return_sequences=2
+                # skip_special_tokens=True,
+                # clean_up_tokenization_spaces=True
+            )
+        comments = [self.tokenizer.decode(generated_id, skip_special_tokens=True) for generated_id in generated_ids]
+        return comments

src/model/text_process.py

@@ -0,0 +1,41 @@
+import re
+
+class TextPreprocess():
+  def __init__(self, teencode_dir="./storage/teencode.txt"):
+    self.get_teencode(teencode_dir)
+
+  def get_teencode(self, teencode_dir):
+    with open(teencode_dir, "r", encoding="utf-8") as f:
+      teencode_original = f.readlines()
+    teencode_json = {}
+    for teencode in teencode_original:
+        key, value = teencode.split("\t")
+        value = value.replace("\n", "")
+        teencode_json[key] = value
+    self.teencode_json = teencode_json
+
+  def teencode_normalize(self, text):
+    text_split = text.split()
+    return " ".join([self.teencode_json.get(txt, txt) for txt in text_split])
+
+  def clean_text(self, text):
+    # Xóa hashtag (dấu #)
+    text = re.sub(r'#\w+', '', text)
+
+    # Xóa liên kết (URL)
+    text = re.sub(r'http\S+', '', text)
+
+    # Xóa các ký tự số
+    text = re.sub(r'\d+', '', text)
+
+    # Xóa ký tự đặc biệt
+    text = re.sub(r'[^\w\s]', '', text)
+
+    text = " ".join(text.split())
+    text = text.lower()
+    return text
+
+  def preprocess(self, text):
+    cleaned_text = self.clean_text(text)
+    cleaned_text = self.teencode_normalize(cleaned_text)
+    return cleaned_text

storage/teencode.txt

@@ -0,0 +1,422 @@
+ctrai	con trai
+khôg	không
+bme	bố mẹ
+cta	chúng ta
+mih	mình
+mqh	mối quan hệ
+cgai	con gái
+nhữg	những
+mng	mọi người
+svtn	sinh viên tình nguyện
+r	rồi
+qtam	quan tâm
+thươg	thương
+qtâm	quan tâm
+chug	chung
+trườg	trường
+thoy	thôi
+đki	đăng ký
+atsm	ảo tưởng sức mạnh
+ạk	ạ
+cv	công việc
+vch	vãi chưởng
+cùg	cùng
+pn	bạn
+pjt	biết
+thjk	thích
+keke	ce ce
+ktra	kiểm tra
+nek	nè
+cgái	con gái
+nthe	như thế
+chúg	chúng
+kái	cái
+tìh	tình
+phòg	phòng
+lòg	lòng
+từg	từng
+rằg	rằng
+sốg	sống
+thuj	thôi
+thuơng	thương
+càg	càng
+đky	đăng ký
+bằg	bằng
+sviên	sinh viên
+ák	á
+đág	đáng
+nvay	như vậy
+nhjeu	nhiều
+xg	xuống
+zồi	rồi
+trag	trang
+zữ	dữ
+atrai	anh trai
+kte	kinh tế
+độg	động
+lmht	liên minh huyền thoại
+gắg	gắng
+đzai	đẹp trai
+thgian	thời gian
+plz	pờ ly
+đồg	đồng
+btrai	bạn trai
+nthê	như thế
+hìhì	hì hì
+vọg	vọng
+hihe	hi he
+đôg	đông
+răg	răng
+thườg	thường
+tcảm	tình cảm
+đứg	đứng
+ksao	không sao
+dz	đẹp trai
+hjxhjx	hix hix
+cmày	chúng mày
+xuốg	xuống
+nkư	như
+lquan	liên quan
+tiếg	tiếng
+hajz	hai
+xih	xinh
+hìh	hình
+thàh	thành
+ngke	nghe
+dzậy	dậy
+teencode	tin cốt
+tnào	thế nào
+tưởg	tưởng
+ctrinh	chương trình
+phog	phong
+hôg	không
+zìa	gì
+kũg	cũng
+ntnao	như thế nào
+trọg	trọng
+nthế	như thế
+năg	năng
+ngđó	người đó
+lquen	làm quen
+riêg	riêng
+ngag	ngang
+hêhê	hê hê
+bnhiu	bao nhiêu
+ngốk	ngốc
+kậu	cậu
+highland	hai lừn
+kqua	kết quả
+htrc	hôm trước
+địh	định
+gđình	gia đinh
+giốg	giống
+csống	cuộc sống
+xug	xùng
+zùi	rồi
+bnhiêu	bao nhiêu
+cbị	chuẩn bị
+kòn	còn
+buôg	buông
+csong	cuộc sống
+chàg	chàng
+chăg	chăng
+ngàh	ngành
+llac	liên lạc
+nkưng	nhưng
+nắg	nắng
+tíh	tính
+khoảg	khoảng
+thík	thích
+ngđo	người đó
+ngkhác	người khác
+thẳg	thẳng
+kảm	cảm
+dàh	dành
+júp	giúp
+lặg	lặng
+vđê	vấn đề
+bbè	bạn bè
+bóg	bóng
+dky	đăng ký
+dòg	dòng
+uốg	uống
+tyêu	tình yêu
+snvv	sinh nhật vui vẻ
+đthoại	điện thoại
+qhe	quan hệ
+cviec	công việc
+tượg	tượng
+qà	quà
+thjc	thích
+nhưq	nhưng
+cđời	cuộc đời
+bthường	bình thường
+zà	già
+đáh	đánh
+xloi	xin lỗi
+zám	dám
+qtrọng	quan trọng
+bìh	bình
+lzi	làm gì
+qhệ	quan hệ
+đhbkhn	đại học bách khoa hà nội
+hajzz	hai
+kủa	của
+lz	làm gì
+đhkhtn	đại học khoa học tự nhiên
+đóg	đóng
+cka	cha
+lgi	làm gì
+nvậy	như vậy
+qả	quả
+đkiện	điều kiện
+nèk	nè
+tlai	tương lai
+bsĩ	bác sĩ
+hkì	học kỳ
+đcsvn	đảng cộng sản việt nam
+vde	vấn đề
+chta	chúng ta
+òy	rồi
+ltinh	linh tinh
+ngyeu	người yêu
+đthoai	điện thoại
+snghĩ	suy nghĩ
+nặg	nặng
+họk	học
+dừg	dừng
+hphúc	hạnh phúc
+hiha	hi ha
+wtâm	quan tâm
+thíck	thích
+chuện	chuyện
+lạh	lạnh
+fây	phây
+ntnày	như thế này
+lúk	lúc
+haj	hai
+ngía	nghía
+mớj	mới
+hsơ	hồ sơ
+ctraj	con trai
+nyêu	người yêu
+điiiiiii	đi
+rồii	rồi
+c	chị
+kih	kinh
+kb	kết bạn
+hixxx	hích
+dthương	dễ thương
+nhiềuuu	nhiều
+ctrình	chương trình
+mìnk	mình
+mjh	mình
+ng	người
+vc	vợ chồng
+uhm	ừm
+thỳ	thì
+nyc	người yêu cũ
+tks	thanks
+nàg	nàng
+thôii	thôi
+đjên	điên
+bgái	bạn gái
+vớii	với
+xink	xinh
+hđộng	hành động
+đhọc	đại học
+mk	mình
+bn	bạn
+thik	thích
+cj	chị
+mn	mọi người
+nguoi	người
+nógn	nóng
+hok	không
+ko	không
+bik	biết
+vs	với
+cx	cũng
+mik	mình
+wtf	what the fuck
+đc	được
+cmt	comment
+ck	chồng
+chk	chồng
+ngta	người ta
+gđ	gia đình
+oh	ồ
+vk	vợ
+ctác	công tác
+sg	sài gòn
+ae	anh em
+ah	à
+ạh	ạ
+rì	gì
+ms	mới
+vn	việt nam
+nhaa	nha
+cũg	cũng
+đag	đang
+ơiii	ơi
+hic	hích
+ace	anh chị em
+àk	à
+uh	ừ
+cmm	con mẹ mày
+cmnr	con mẹ nó rồi
+ơiiii	ơi
+hnay	hôm nay
+ukm	ừm
+tq	trung quốc
+ctr	chương trình
+đii	đi
+nch	nói chuyện
+trieu	triệu
+hahah	ha ha
+nta	người ta
+ngèo	nghèo
+kêh	kênh
+ak	à
+ad	admin
+j	gì
+ny	người yêu
+dc	được
+qc	quảng cáo
+baoh	bao giờ
+zui	vui
+zẻ	vẻ
+tym	tim
+aye	anh yêu em
+eya	em yêu anh
+fb	facebook
+insta	instagram
+z	vậy
+thich	thích
+vcl	vờ cờ lờ
+đt	điện thoại
+acc	account
+lol	lồn
+loz	lồn
+lozz	lồn
+trc	trước
+chs	chẳng hiểu sao
+đhs	đéo hiểu sao
+qá	quá
+ntn	như thế nào
+wá	quá
+zậy	vậy
+zô	vô
+ytb	youtube
+vđ	vãi đái
+vchg	vãi chưởng
+sml	sấp mặt lờ
+xl	xin lỗi
+cmn	con mẹ nó
+face	facebook
+hjhj	hi hi
+vv	vui vẻ
+ns	nói
+iu	yêu
+vcđ	vãi cải đái
+in4	info
+qq	quằn què
+sub	subcribe
+kh	không
+zạ	vậy
+oy	rồi
+jo	giờ
+clmm	cái lồn mẹ mày
+bsvv	buổi sáng vui vẻ
+troai	trai
+wa	quá
+hjx	hix
+e	em
+ik	đi
+ji	gì
+ce	chị em
+lm	làm
+đz	đẹp giai
+sr	sorry
+ib	inbox
+hoy	thôi
+đbh	đéo bao giờ
+k	không
+vd	ví dụ
+a	anh
+cũng z	cũng vậy
+z là	vậy là
+unf	unfriend
+my fen	my friend
+fen	friend
+cty	công ty
+on lai	online
+u hai ba	u23
+kô	không
+đtqg	đội tuyển quốc gia
+hqua	hôm qua
+xog	xong
+uh	ừ
+uk	ừ
+nhoé	nhé
+biet	biết
+quí	quý
+stk	số tài khoản
+hong kong	hồng kông
+đươc	được
+nghành	ngành
+nvqs	nghĩa vụ quân sự
+ngừoi	người
+trog	trong
+tgian	thời gian
+biêt	biết
+fải	phải
+nguời	người
+tđn	thế đéo nào
+bth	bình thường
+vcđ	vãi cả đái
+tgdd	thế giới di động
+khg	không
+nhưg	nhưng
+thpt	trung học phổ thông
+thằg	thằng
+đuợc	được
+dc	được
+đc	được
+ah	à
+àh	à
+ku	cu
+thým	thím
+onl	online
+zô	dô
+zú	vú
+cmnd	chứng minh nhân dân
+sđt	số điện thoại
+klq	không liên quan
+t	mình
+vại	vậy
+v	vậy
+z	vậy
+cám ơn	cảm ơn
+mún	muốn
+thặc	thật
+gấc	rất
+zị	vậy
+cf	cà phê
+kout	cao
+thưn	thương
+rùi	rồi
+thui	thôi
+hong	không
+khum	không
+chòi	trời
+b	bạn
+lém	lắm
+c	chị
+gòi	rồi
+zị	vậy
+chài	trời
+lun	luôn
+=	bằng
+đsung	đúng
+tr	trời

utils/config.py

@@ -0,0 +1,18 @@
+import yaml
+
+class Config():
+    def __init__(self, dir) -> None:
+        """Init function for Config class
+
+        Args:
+            dir (String): path to config file
+        """
+        self.dir = dir
+    
+    def __get_config__(self):
+        """get config
+
+        Returns:
+            yaml: Yaml config
+        """
+        return yaml.load((open(self.dir)), Loader=yaml.SafeLoader)

utils/get_logging.py

@@ -0,0 +1,18 @@
+import pytz
+import logging
+import datetime
+
+def get_logging(log_dir):
+    """Get logging function
+
+    Args:
+        log_dir (String): location directory (relative path)
+
+    Returns:
+        logging: custom logging
+    """
+    
+    logging.Formatter.converter = logging.Formatter.converter = lambda *args: datetime.datetime.now(tz=pytz.timezone('Asia/Ho_Chi_Minh')).timetuple()
+    logging.basicConfig(filename=log_dir, level=logging.INFO, format = '[%(asctime)s] - [%(levelname)s] - [%(funcName)s] - [%(lineno)d] - %(message)s',
+                    filemode="w")
+    return logging