demo

README.md

@@ -1,12 +1,13 @@
 ---
-title: Test3
-emoji: 📈
+title: Kobedemo
+emoji: 🐢
 colorFrom: red
-colorTo: purple
+colorTo: red
 sdk: streamlit
 sdk_version: 1.17.0
 app_file: app.py
 pinned: false
+python_version: 3.8.12
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,2 @@
+import streamlit as st
+st.write('You have selected:')

app1.py

@@ -0,0 +1,198 @@
+"""
+import argparse
+import random
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from dataset import KobeDataModule
+from model import KobeModel
+from options import add_args, add_options
+from typing import List
+import warnings
+import logging
+from transformers.models.bert.tokenization_bert import BertTokenizer
+import sentencepiece as spm
+
+
+logging.getLogger("lightning").setLevel(logging.ERROR)
+def fxn():
+    warnings.warn("deprecated", DeprecationWarning)
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    fxn()
+
+
+    
+parser = argparse.ArgumentParser()
+add_options(parser)
+args = parser.parse_args()
+args.vocab_file= "bert-base-chinese" 
+args.cond_vocab_file = "./vocab.cond.model"
+add_args(args)
+ 
+
+    
+#model = KobeModel(args)
+
+
+
+
+class Example:
+    title_token_ids: List[int]
+    condition_token_ids: List[int]
+    fact_token_ids: List[int]
+
+    def __init__(self, title_token_ids:List[int], condition_token_ids: List[int],  fact_token_ids: List[int]):
+        self.title_token_ids = title_token_ids
+        self.condition_token_ids = condition_token_ids
+        self.fact_token_ids = fact_token_ids
+
+text_tokenizer = BertTokenizer.from_pretrained(args.vocab_file)
+cond_tokenizer = spm.SentencePieceProcessor()
+cond_tokenizer.Load(args.cond_vocab_file)
+
+
+
+
+
+#model = model.load_from_checkpoint("/root/kobe-v2/1ja19m5t/checkpoints/epoch=19-step=66080.ckpt", args=args)
+#model = model.load_from_checkpoint("/root/kobe-v2/37ht1cvz/checkpoints/epoch=11-step=396384.ckpt", args=args)
+
+trainer = pl.Trainer(accelerator='gpu', devices=1, max_epochs=-1)
+"""
+
+
+import streamlit as st
+
+st.write("Most appearing words including stopwords")
+
+"""
+choice = st.selectbox(
+
+    'Select the items you want?',
+
+    ('Pen','Pencil','Eraser','Sharpener','Notebook'))
+
+
+input1 = st.selectbox(
+    'please choose:',
+    ("1.天猫直发百诺碳纤维三脚架单反相机三角架c2690tb1摄影脚架",
+    "2.dacom飞鱼p10运动型跑步蓝牙耳机入耳式头戴挂耳塞式7级防水苹果安卓手机通用可接听电话音乐篮",
+    "3.coach蔻驰贝壳包pvc单肩手提斜挎大号贝壳包女包5828",
+    "4.highcook韩库韩国进口蓝宝石近无烟炒锅家用不粘锅电磁炉炒菜锅",
+    "5.欧式复古亚麻布料沙发面料定做飘窗垫窗台垫榻榻米垫抱枕diy布艺",
+    "6.飞利浦电动剃须刀sp9851充电式带多功能理容配件和智能清洁系统",
+    "7.不锈钢牛排刀叉西餐餐具全套筷子刀叉勺三件套欧式加厚24件礼盒装",
+    "8.香百年汽车香水挂件车载香水香薰悬挂吊坠车用车内装饰挂饰精油",
+    "9.迪士尼小学生书包儿童男孩13一46年级美国队长男童减负12周岁男",
+    "10.半饱良味潮汕猪肉脯宅人食堂潮汕小吃特产碳烤猪肉干120g"))
+
+st.write('You selected:', input1)
+
+
+title =""
+fact = ""
+
+if input1==1:
+    title= "天猫直发百诺碳纤维三脚架单反相机三角架c2690tb1摄影脚架"
+    fact = "太字节Terabyte，计算机存储容量单位，也常用TB来表示。百诺公司创建于1996年，早期与日本合作，后通过自身技术创新与努力，逐渐在国内外抑尘设备行业赢得一席单反就是指单镜头反光，即SLRSingleLensReflex，单反相机就是拥有这个功能的相机。技巧拍摄往往都离不开三脚架的帮助，如夜景拍摄、微距拍摄等方面。"
+    
+
+if input1==2:
+    title="dacom飞鱼p10运动型跑步蓝牙耳机入耳式头戴挂耳塞式7级防水苹果安卓手机通用可接听电话音乐篮牙"
+    fact = "移动电话，或称为无线电话，通常称为手机，原本只是一种通讯工具，早期又有大哥大的俗称，是可以在较广范围运动型是德国精神病学家克雷奇默提出的身体类型之一。跑步，是指陆生动物使用足部，移动最快捷的方法。蓝牙耳机就是将蓝牙技术应用在免持耳机上，让使用者可以免除恼人电线的牵绊，自在地以各种方式轻松通话。"
+
+
+if input1==3:
+    title="coach蔻驰贝壳包pvc单肩手提斜挎大号贝壳包女包5828"
+    fact = "聚氯乙烯，英文简称PVCPolyvinylchloride，是氯乙烯单体vinylchloridem女包，这个名词是箱包的性别分类衍生词。贝壳包beikebao女士包种类的一种，因为其外形酷似贝壳的外形而得名。蔻驰为美国经典皮件品牌COACH，一像以简洁、耐用的风格特色赢得消费者的喜爱。"
+
+
+if input1==4:
+    title= "highcook韩库韩国进口蓝宝石近无烟炒锅家用不粘锅电磁炉炒菜锅"
+    fact = "蓝宝石，是刚玉宝石中除红色的红宝石之外，其它颜色刚玉宝石的通称，主要成分是氧化铝Al2O3。电磁炉又称为电磁灶，1957年第一台家用电磁炉诞生于德国。家用是汉语词汇，出自管子权修，解释为家庭日常使用的。不粘锅即做饭不会粘锅底的锅，是因为锅底采用了不粘涂层，常见的、不粘性能最好的有特���龙涂层和陶瓷涂层。" 
+
+
+if input1==5:
+    title= "欧式复古亚麻布料沙发面料定做飘窗垫窗台垫榻榻米垫抱枕diy布艺"
+    fact = "沙发是个外来词，根据英语单词sofa音译而来。面料就是用来制作服装的材料。飘窗垫，就是放在飘窗的台面上的垫子。复古与怀旧，有时候很难区分。" 
+
+
+if input1==6:
+    title= "飞利浦电动剃须刀sp9851充电式带多功能理容配件和智能清洁系统"
+    fact = "能够完成一种或者几种生理功能的多个器官按照一定的次序组合在一起的结构叫做系统。配件，指装配机械的零件或部件；也指损坏后重新安装上的零件或部件。清洁是由奥利维耶阿萨亚斯执导，张曼玉、尼克诺尔蒂主演的剧情片，于2004年9月1日在法国上映。飞利浦，1891年成立于荷兰，主要生产照明、家庭电器、医疗系统方面的产品。"
+
+if input1==7:
+    title= "不锈钢牛排刀叉西餐餐具全套筷子刀叉勺三件套欧式加厚24件礼盒装"
+    fact = "西餐餐具具体有大盘子、小盘子、浅碟、深碟、吃沙拉用的叉子、叉肉用的叉子、喝汤用的汤匙、吃甜点用的汤匙三件咳嗽,贫穷与爱情触不到的恋人...不锈钢指耐空气、蒸汽、水等弱腐蚀介质和酸、碱、盐等化学浸蚀性介质腐蚀的钢，又称不锈耐酸钢。2，4D丁酯，无色油状液体。"
+
+if input1==8:
+    title= "香百年汽车香水挂件车载香水香薰悬挂吊坠车用车内装饰挂饰精油"
+    fact = "悬挂系统是汽车的车架与车桥或车轮之间的一切传力连接装置的总称，其作用是传递作用在车轮和车架之间的力和吊坠，一种首饰，配戴在脖子上的饰品，多为金属制，特别是不锈钢制和银制，也有矿石、水晶、玉石等制的，主汽车香水AutoPerfume是一种混合了香精油、固定剂与酒精的液体，用来让汽车车内拥有持久且悦人的精油是从植物的花、叶、茎、根或果实中，通过水蒸气蒸馏法、挤压法、冷浸法或溶剂提取法提炼萃取的挥发性芳"
+
+if input1==9:
+    title = "迪士尼小学生书包儿童男孩13一46年级美国队长男童减负12周岁男"
+    fact= "美国队长是每一男孩心中的英雄人物，迪士尼美国队长款的小学生书包，按照美国队长的防护盾牌设计，泛着丝丝银光，帅气有型，而且还有很多英雄款式哦脊背处采用柔软舒适的脊椎防护设计，减轻孩子的背部压力。而且前方盾牌还能拆卸下来，当做斜挎包使用，满足淘气小男孩的英雄梦。"
+   
+if input1==10:
+    title = "半饱良味潮汕猪肉脯宅人食堂潮汕小吃特产碳烤猪肉干120g"
+    fact = "潮汕，不是潮州潮州一词始于隋文帝开皇十年，距今不到两千年。猪肉脯是一种用猪肉经腌制、烘烤的片状肉制品，食用方便、制作考究、美味可口、耐贮藏和便于运输的中式传统发达国家都有全国统一的急救电话号码。特产指某地特有的或特别著名的产品，有文化内涵或历史，亦指只有在某地才生产的一种产品。"
+    
+
+input2 = st.selectbox(
+    'please choose category:',
+    ("1: 家庭主妇",
+    "2: 烹饪达人",
+    "3: 买鞋控",
+    "4: 数码达人",
+    "5: 吃货",
+    "6: 爱包人",
+    "7: 高富帅",
+    ))
+
+st.write('You selected:', input2)
+aspect = "<"+str(input2)+">"
+
+
+
+input3 = st.selectbox(
+    'please choose aspect:',
+    ("1: appearance",
+    "2: texture",
+    "3: function"))
+
+st.write('You selected:', input3)
+
+cond = ""
+if input3==1:
+    cond="<a>"
+if input3==2:
+    cond="<b>"
+if input3==3:
+    cond="<c>"
+
+#cond = cond+" "+aspect
+cond = aspect+" "+cond
+#print(title)
+#print(fact)
+#print(cond)
+tokenizer = text_tokenizer
+title_token_ids=tokenizer.encode(title, add_special_tokens=False)
+condition_token_ids=cond_tokenizer.EncodeAsIds(cond)
+fact_token_ids=tokenizer.encode(fact, add_special_tokens=False)
+
+e = Example(title_token_ids, condition_token_ids, fact_token_ids)
+
+
+dm = KobeDataModule(
+        [e],
+        args.text_vocab_path,
+        args.max_seq_len,
+        1,
+        1,
+    )
+for d in dm.test_dataloader():
+    st.write("result:")
+    st.write(''.join(model.test_step(d ,1)).replace(" ",""))
+""" 

dataset.py

@@ -0,0 +1,271 @@
+import glob
+from dataclasses import dataclass
+from typing import List
+
+import pytorch_lightning as pl
+import sentencepiece as spm
+import torch
+from torch.functional import Tensor
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data.dataloader import DataLoader
+
+
+
+@dataclass
+class Example:
+    title_token_ids: List[int]
+    description_token_ids: List[int]
+    condition_token_ids: List[int]
+    fact_token_ids: List[int]
+    description: str
+    title: str
+
+
+@dataclass
+class TensorDict:
+    def detach(self):
+        detached_dict = {
+            field: getattr(self, field).detach()
+            if isinstance(getattr(self, field), torch.Tensor)
+            else getattr(self, field)
+            for field in self.__dataclass_fields__
+        }
+        return self.__class__(**detached_dict)
+
+    def cpu(self):
+        detached_dict = {
+            field: getattr(self, field).cpu()
+            if isinstance(getattr(self, field), torch.Tensor)
+            else getattr(self, field)
+            for field in self.__dataclass_fields__
+        }
+        return self.__class__(**detached_dict)
+
+
+@dataclass
+class Batched(TensorDict):
+    # Source
+    title_token_ids: torch.Tensor
+    title_token_ids_mask: torch.Tensor
+    # Attribute Fusion
+    cond_title_token_ids: torch.Tensor
+    cond_title_token_ids_mask: torch.Tensor
+    # Knowledge Incorporation
+    fact_token_ids: torch.Tensor
+    fact_token_ids_mask: torch.Tensor
+    title_fact_token_ids: torch.Tensor
+    title_fact_token_ids_mask: torch.Tensor
+    # Attribute Fusion + Knowledge Incorporation
+    cond_title_fact_token_ids: torch.Tensor
+    cond_title_fact_token_ids_mask: torch.Tensor
+    # Target
+    #description_token_ids: torch.Tensor
+    #description_token_ids_mask: torch.Tensor
+    #descriptions: List[str]
+    #titles: List[str]
+
+
+
+@dataclass
+class EncodedBatch(TensorDict):
+    context_encodings: torch.Tensor
+    context_encodings_mask: torch.Tensor
+
+
+@dataclass
+class DecodedBatch:
+    loss: float
+    acc: float
+    generated: List[str]
+    descriptions: List[str]
+    titles: List[str]
+
+
+def from_processed(url: str, train=False):
+    urls = sorted(glob.glob(url))
+    def my_split_by_worker(urls):
+        wi = torch.utils.data.get_worker_info()
+        if wi is None:
+            return urls
+        else:
+            return urls[wi.id::wi.num_workers]
+    def my_split_by_node(urls):
+        node_id, node_count = torch.distributed.get_rank(), torch.distributed.get_world_size()
+        return urls[node_id::node_count]
+    if train:
+
+        return (
+            wds.WebDataset(urls)
+            #wds.WebDataset(urls,nodesplitter=my_split_by_node)
+            #wds.WebDataset(urls,nodesplitter=wds.split_by_node)
+            .shuffle(size=10000000, initial=100000)
+            .decode()
+            .map(lambda d: Example(**d["json"]))
+        ) 
+    else:
+        print(list(wds.WebDataset(url).decode().map(lambda d: Example(**d["json"])))[0])
+        sys.exit()
+        return list(wds.WebDataset(url).decode().map(lambda d: Example(**d["json"])))
+        #return list(wds.WebDataset(urls, nodesplitter=my_split_by_node).decode().map(lambda d: Example(**d["json"])))
+        #return list(wds.WebDataset(urls, nodesplitter=wds.split_by_node).decode().map(lambda d: Example(**d["json"])))
+
+
+def get_collate_fn(text_vocab_size: int, max_seq_length: int):
+    def collate_fn(examples: List[Example]) -> Batched:
+        from kobe.data.vocab import BOS_ID, EOS_ID
+
+        title_token_ids = pad_sequence(
+            [
+                torch.tensor(
+                    [BOS_ID] + e.title_token_ids[: max_seq_length - 2] + [EOS_ID]
+                )
+                for e in examples
+            ]
+        )
+        fact_token_ids = pad_sequence(
+            [
+                torch.tensor(
+                    [BOS_ID] + e.fact_token_ids[: max_seq_length - 2] + [EOS_ID]
+                )
+                for e in examples
+            ]
+        )
+        """
+        description_token_ids = pad_sequence(
+            [
+                torch.tensor(
+                    [BOS_ID] + e.description_token_ids[: max_seq_length - 2] + [EOS_ID]
+                )
+                for e in examples
+            ]
+        )
+        """
+        cond_title_token_ids = pad_sequence(
+            [
+                torch.tensor(
+                    (
+                        [BOS_ID]
+                        + [
+                            cond_id + text_vocab_size
+                            for cond_id in e.condition_token_ids
+                        ]
+                        + e.title_token_ids
+                    )[: max_seq_length - 1]
+                    + [EOS_ID]
+                )
+                for e in examples
+            ]
+        )
+        title_fact_token_ids = pad_sequence(
+            [
+                torch.tensor(
+                    ([BOS_ID] + e.title_token_ids + [EOS_ID] + e.fact_token_ids)[
+                        : max_seq_length - 1
+                    ]
+                    + [EOS_ID]
+                )
+                for e in examples
+            ]
+        )
+        cond_title_fact_token_ids = pad_sequence(
+            [
+                torch.tensor(
+                    (
+                        [BOS_ID]
+                        + [
+                            cond_id + text_vocab_size
+                            for cond_id in e.condition_token_ids
+                        ]
+                        + e.title_token_ids
+                        + [EOS_ID]
+                        + e.fact_token_ids
+                    )[: max_seq_length - 1]
+                    + [EOS_ID]
+                )
+                for e in examples
+            ]
+        )
+        #descriptions = [e.description for e in examples]
+        #titles = [e.title for e in examples]
+        return Batched(
+            title_token_ids=title_token_ids,
+            title_token_ids_mask=(title_token_ids == 0).T,
+            fact_token_ids=fact_token_ids,
+            fact_token_ids_mask=(fact_token_ids == 0).T,
+            cond_title_token_ids=cond_title_token_ids,
+            cond_title_token_ids_mask=(cond_title_token_ids == 0).T,
+            title_fact_token_ids=title_fact_token_ids,
+            title_fact_token_ids_mask=(title_fact_token_ids == 0).T,
+            cond_title_fact_token_ids=cond_title_fact_token_ids,
+            cond_title_fact_token_ids_mask=(cond_title_fact_token_ids == 0).T,
+            #description_token_ids="",
+            #description_token_ids_mask=(description_token_ids == 0).T,
+            #descriptions="",
+            #titles="",
+        )
+
+    return collate_fn
+
+
+class KobeDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        test_data: str,
+        vocab_path: str,
+        max_seq_length: int,
+        batch_size: int,
+        num_workers: int,
+    ):
+        super().__init__()
+        self.test_data = test_data
+        self.max_seq_length = max_seq_length
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.text_vocab_size = helpers.get_bert_vocab_size(vocab_path)
+
+
+    """
+    def train_dataloader(self):
+        return DataLoader(
+            self.train,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=get_collate_fn(self.text_vocab_size, self.max_seq_length),
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.valid,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=get_collate_fn(self.text_vocab_size, self.max_seq_length),
+        )
+    """
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_data,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            collate_fn=get_collate_fn(self.text_vocab_size, self.max_seq_length),
+        )
+
+
+if __name__ == "__main__":
+    dm = KobeDataModule(
+        train_data="saved/processed/train-*.tar",
+        valid_data="saved/processed/valid.tar",
+        test_data="saved/processed/test.tar",
+        vocab_path="bert-base-chinese",
+        max_seq_length=512,
+        batch_size=32,
+        num_workers=8,
+    )
+    dm.setup("test")
+    max_len = 0
+    from tqdm import tqdm
+
+    tqdm_iter = tqdm(dm.test_dataloader())
+    for batch in tqdm_iter:
+        max_len = max(max_len, batch.cond_title_fact_token_ids.shape[0])
+        max_len = max(max_len, batch.description_token_ids.shape[0])
+        tqdm_iter.set_description(f"max len = {max_len}")

helper.py

@@ -0,0 +1,79 @@
+import sentencepiece as spm
+import torch
+import torch.nn.functional as F
+from transformers.models.bert.tokenization_bert import BertTokenizer
+
+BASELINE = "baseline"
+KOBE_ATTRIBUTE = "kobe-attr"
+KOBE_KNOWLEDGE = "kobe-know"
+KOBE_FULL = "kobe-full"
+
+
+def get_bert_vocab_size(vocab_path: str) -> int:
+    tokenizer = BertTokenizer.from_pretrained(vocab_path)
+    return tokenizer.vocab_size
+
+
+def get_vocab_size(vocab_path: str) -> int:
+    tokenizer = spm.SentencePieceProcessor()
+    tokenizer.Load(vocab_path)
+    return len(tokenizer)
+
+
+
+# Metrics
+def accuracy(logits: torch.Tensor, targets: torch.Tensor) -> float:
+    assert logits.dim() == 2
+    assert targets.dim() == 1
+    pred = logits.argmax(dim=1)
+    return (pred == targets).sum().item() / targets.shape[0]
+
+
+def top_k_top_p_sampling(
+    logits, top_k=0, top_p=0.0, temperature=1, filter_value=-float("Inf")
+) -> int:
+    """Sample from a filtered distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (vocabulary size)
+        top_k >0: keep only top k tokens with highest probability (top-k filtering).
+        top_p >0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+    """
+    logits /= temperature
+    assert (
+        logits.dim() == 1
+    )  # batch size 1 for now - could be updated for more but the code would be less clear
+    top_k = min(top_k, logits.size(-1))  # Safety check
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+
+    # Sample from the filtered distribution
+    probabilities = F.softmax(logits, dim=-1)
+    next_token = torch.multinomial(probabilities, 1)
+
+    return int(next_token.item())
+
+
+def diversity(tokenized_lines, n=4) -> int:
+    """Defined as the unique number of ngrams generated on the test set."""
+    n_grams_all = []
+    for line in tokenized_lines:
+        n_grams = list(zip(*[line[i:] for i in range(n)]))
+        n_grams_all += n_grams
+
+    return len(set(n_grams_all))

model.py

@@ -0,0 +1,131 @@
+from typing import List, Tuple
+
+import numpy as np
+import pytorch_lightning as pl
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+
+#from sacrebleu.metrics.bleu import BLEU, _get_tokenizer
+from torch import optim
+from torch.nn.init import xavier_uniform_
+from transformers.models.bert.tokenization_bert import BertTokenizer
+
+#import wandb
+from dataset import Batched, DecodedBatch
+#from models.scheduler import WarmupDecayLR
+from transformer import Decoder, Encoder
+#from kobe.utils import helpers
+
+
+class KobeModel(pl.LightningModule):
+    def __init__(self, args):
+        super(KobeModel, self).__init__()
+
+        self.encoder = Encoder(
+            vocab_size=args.text_vocab_size + args.cond_vocab_size,
+            max_seq_len=args.max_seq_len,
+            d_model=args.d_model,
+            nhead=args.nhead,
+            num_layers=args.num_encoder_layers,
+            dropout=args.dropout,
+            mode=args.mode,
+        )
+        self.decoder = Decoder(
+            vocab_size=args.text_vocab_size,
+            max_seq_len=args.max_seq_len,
+            d_model=args.d_model,
+            nhead=args.nhead,
+            num_layers=args.num_decoder_layers,
+            dropout=args.dropout,
+        )
+        self.lr = args.lr
+        self.d_model = args.d_model
+        self.loss = nn.CrossEntropyLoss(
+            reduction="mean", ignore_index=0, label_smoothing=0.1
+        )
+        self._reset_parameters()
+
+        self.decoding_strategy = args.decoding_strategy
+        self.vocab = BertTokenizer.from_pretrained(args.text_vocab_path)
+        #self.bleu = BLEU(tokenize=args.tokenize)
+        #self.sacre_tokenizer = _get_tokenizer(args.tokenize)()
+        #self.bert_scorer = BERTScorer(lang=args.tokenize, rescale_with_baseline=True)
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                xavier_uniform_(p)
+
+    def _tokenwise_loss_acc(
+        self, logits: torch.Tensor, batch: Batched
+    ) -> Tuple[torch.Tensor, float]:
+        unmask = ~batch.description_token_ids_mask.T[1:]
+        unmasked_logits = logits[unmask]
+        unmasked_targets = batch.description_token_ids[1:][unmask]
+        #acc = helpers.accuracy(unmasked_logits, unmasked_targets)
+        return self.loss(logits.transpose(1, 2), batch.description_token_ids[1:]), 1
+
+    def training_step(self, batch: Batched, batch_idx: int):
+        encoded = self.encoder.forward(batch)
+        logits = self.decoder.forward(batch, encoded)
+        loss, acc = self._tokenwise_loss_acc(logits, batch)
+        self.lr_schedulers().step()
+        self.log("train/loss", loss.item())
+        self.log("train/acc", acc)
+        return loss
+
+    def _shared_eval_step(self, batch: Batched, batch_idx: int) -> DecodedBatch:
+        encoded = self.encoder.forward(batch)
+        #logits = self.decoder.forward(batch, encoded)
+        #loss, acc = self._tokenwise_loss_acc(logits, batch)
+
+        preds = self.decoder.predict(
+            encoded_batch=encoded, decoding_strategy=self.decoding_strategy
+        )
+        generated = self.vocab.batch_decode(preds.T.tolist(), skip_special_tokens=True)
+        #print(generated)
+
+        return generated
+        return DecodedBatch(
+            loss=loss.item(),
+            acc=acc,
+            generated=generated,
+            descriptions=batch.descriptions,
+            titles=batch.titles,
+        )
+
+    def validation_step(self, batch, batch_idx):
+        return self._shared_eval_step(batch, batch_idx)
+
+    def test_step(self, batch, batch_idx, dataloader_idx=0):
+        return self._shared_eval_step(batch, batch_idx)
+
+    def _shared_epoch_end(self, outputs: List[DecodedBatch], prefix):
+        loss = np.mean([o.loss for o in outputs])
+        acc = np.mean([o.acc for o in outputs])
+        self.log(f"{prefix}/loss", loss)
+        self.log(f"{prefix}/acc", acc)
+        print(outputs)
+
+        generated = [g for o in outputs for g in o.generated]
+        references = [r for o in outputs for r in o.descriptions]
+        titles = [r for o in outputs for r in o.titles]
+
+
+        # Examples
+        columns = ["Generated", "Reference"]
+        data = list(zip(generated[:256:16], references[:256:16]))
+        table = wandb.Table(data=data, columns=columns)
+        self.logger.experiment.log({f"examples/{prefix}": table})
+
+    def validation_epoch_end(self, outputs):
+        self._shared_epoch_end(outputs, "val")
+
+    def test_epoch_end(self, outputs):
+        self._shared_epoch_end(outputs, "test")
+
+    def configure_optimizers(self):
+        optimizer = optim.AdamW(self.parameters(), lr=self.lr, betas=(0.9, 0.98))
+        #scheduler = WarmupDecayLR(optimizer, warmup_steps=10000, d_model=self.d_model)
+        return [optimizer]

options.py

@@ -0,0 +1,47 @@
+from argparse import ArgumentParser, Namespace
+
+import helpers
+
+
+def add_options(parser: ArgumentParser):
+    # fmt: off
+    # Dataset
+    parser.add_argument("--train-data", default="saved/processed/train-*.tar", type=str)
+    parser.add_argument("--valid-data", default="saved/processed/valid.tar", type=str)
+    parser.add_argument("--test-data", default="saved/processed/test.tar", type=str)
+    parser.add_argument("--text-vocab-path", default="bert-base-chinese", type=str, help="BertTokenizer used to preprocess the corpus")
+    parser.add_argument("--cond-vocab-path", default="./vocab.cond.model", type=str)
+    parser.add_argument("--num-workers", default=8, help="Number of data loaders", type=int)
+    parser.add_argument("--tokenize", default="zh", help="Tokenization method used to compute sacrebleu, diversity, and BERTScore, defaulted to Chinese", type=str)
+
+    # Model
+    parser.add_argument("--d-model", default=512, type=int)
+    parser.add_argument("--nhead", default=8, type=int)
+    parser.add_argument("--num-encoder-layers", default=6, type=int)
+    parser.add_argument("--num-decoder-layers", default=6, type=int)
+    parser.add_argument("--max-seq-len", default=256, type=int)
+    parser.add_argument("--mode", default="baseline", type=str, choices=[
+        helpers.BASELINE, helpers.KOBE_ATTRIBUTE, helpers.KOBE_KNOWLEDGE, helpers.KOBE_FULL])
+
+    # Training
+    parser.add_argument("--name", default="exp", type=str, help="expeirment name")
+    parser.add_argument("--gpu", default=1, type=int)
+    parser.add_argument("--grad-clip", default=1.0, type=float, help="clip threshold of gradients")
+    parser.add_argument("--epochs", default=30, type=int, help="number of epochs to train")
+    parser.add_argument("--patience", default=10, type=int, help="early stopping patience")
+    parser.add_argument("--lr", default=1, type=float, help="learning rate")
+    parser.add_argument("--dropout", default=0.1, type=float, help="dropout rate")
+    parser.add_argument("--batch-size", default=64, type=int)
+    parser.add_argument("--seed", default=42, type=int)
+
+    # Evaluation
+    parser.add_argument("--test", action="store_true", help="only do evaluation")
+    parser.add_argument("--load-file", required=False, type=str, help="path to the checkpoint (.ckpt) for evaluation")
+    parser.add_argument("--decoding-strategy", default="greedy", type=str, choices=["greedy", "nucleus"], help="Whether to use greedy decoding or nucleus sampling (https://arxiv.org/abs/1904.09751)")
+
+    # fmt: on
+
+
+def add_args(args: Namespace):
+    args.text_vocab_size = helpers.get_bert_vocab_size(args.text_vocab_path)
+    args.cond_vocab_size = helpers.get_vocab_size(args.cond_vocab_path)

requirements.txt

@@ -0,0 +1,6 @@
+torch==1.10.0
+transformers==4.25.1
+sentencepiece
+pytorch-lightning==1.6.4  
+  
+

transformer.py

@@ -0,0 +1,222 @@
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from cached_property import cached_property
+from torch.nn.modules.transformer import (
+    TransformerDecoder,
+    TransformerDecoderLayer,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+
+from dataset import Batched, EncodedBatch
+from vocab import BOS_ID, EOS_ID, PAD_ID
+import helper
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, dropout, dim, max_len=5000):
+        """
+        initialization of required variables and functions
+        :param dropout: dropout probability
+        :param dim: hidden size
+        :param max_len: maximum length
+        """
+        super(PositionalEncoding, self).__init__()
+        # positional encoding initialization
+        pe = torch.zeros(max_len, dim)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        # term to divide
+        div_term = torch.exp(
+            (torch.arange(0, dim, 2, dtype=torch.float) * -(math.log(10000.0) / dim))
+        )
+        # sinusoidal positional encoding
+        pe[:, 0::2] = torch.sin(position.float() * div_term)
+        pe[:, 1::2] = torch.cos(position.float() * div_term)
+        pe = pe.unsqueeze(1)
+        self.register_buffer("pe", pe)
+        self.dropout = nn.Dropout(p=dropout)
+        self.dim = dim
+
+    def forward(self, emb):
+        """
+        create positional encoding
+        :param emb: word embedding
+        :param step: step for decoding in inference
+        :return: positional encoding representation
+        """
+        emb *= math.sqrt(self.dim)
+        emb = emb + self.pe[: emb.size(0)]  # [len, batch, size]
+        emb = self.dropout(emb)
+        return emb
+
+
+class Encoder(nn.Module):
+    @staticmethod
+    def from_args(args) -> "Encoder":
+        return Encoder(
+            args.text_vocab_size + args.cond_vocab_size,
+            args.max_seq_len,
+            args.d_model,
+            args.nhead,
+            args.num_encoder_layers,
+            args.dropout,
+            args.mode,
+        )
+
+    def __init__(
+        self,
+        vocab_size: int,
+        max_seq_len: int,
+        d_model: int,
+        nhead: int,
+        num_layers: int,
+        dropout: float,
+        mode: str,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.max_seq_len = max_seq_len
+        self.input_embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_encoder = PositionalEncoding(dropout, d_model)
+        encoder_layer = TransformerEncoderLayer(
+            d_model, nhead, d_model * 4, dropout, norm_first=True
+        )
+        self.encoder = TransformerEncoder(
+            encoder_layer, num_layers, nn.LayerNorm(d_model)
+        )
+        self.mode = mode
+
+    @cached_property
+    def device(self):
+        return list(self.parameters())[0].device
+
+    def forward(self, batched: Batched) -> EncodedBatch:
+        src, src_key_padding_mask = Encoder._get_input(batched, self.mode)
+        src = self.input_embedding(src)
+        src = self.pos_encoder(src)
+        token_encodings = self.encoder.forward(
+            src=src, src_key_padding_mask=src_key_padding_mask
+        )
+        return EncodedBatch(
+            context_encodings=token_encodings,
+            context_encodings_mask=src_key_padding_mask,
+        )
+
+    @staticmethod
+    def _get_input(batched: Batched, mode: str) -> Tuple[torch.Tensor, torch.Tensor]:
+        return {
+            helpers.BASELINE: (batched.title_token_ids, batched.title_token_ids_mask),
+            helpers.KOBE_ATTRIBUTE: (
+                batched.cond_title_token_ids,
+                batched.cond_title_token_ids_mask,
+            ),
+            helpers.KOBE_KNOWLEDGE: (
+                batched.title_fact_token_ids,
+                batched.title_fact_token_ids_mask,
+            ),
+            helpers.KOBE_FULL: (
+                batched.cond_title_fact_token_ids,
+                batched.cond_title_fact_token_ids_mask,
+            ),
+        }[mode]
+
+
+class Decoder(nn.Module):
+    @staticmethod
+    def from_args(args) -> "Decoder":
+        return Decoder(
+            args.text_vocab_size,
+            args.max_seq_len,
+            args.d_model,
+            args.nhead,
+            args.num_encoder_layers,
+            args.dropout,
+        )
+
+    def __init__(
+        self,
+        vocab_size: int,
+        max_seq_len: int,
+        d_model: int,
+        nhead: int,
+        num_layers: int,
+        dropout: float,
+    ):
+        super(Decoder, self).__init__()
+        self.max_seq_len = max_seq_len
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_encoder = PositionalEncoding(dropout, d_model)
+        decoder_layer = TransformerDecoderLayer(
+            d_model, nhead, 4 * d_model, dropout, norm_first=True
+        )
+        self.decoder = TransformerDecoder(
+            decoder_layer, num_layers, nn.LayerNorm(d_model)
+        )
+        self.output = nn.Linear(d_model, vocab_size)
+
+    def forward(self, batch: Batched, encoded_batch: EncodedBatch) -> torch.Tensor:
+        tgt = self.embedding(batch.description_token_ids[:-1])
+        tgt = self.pos_encoder(tgt)
+        tgt_mask = Decoder.generate_square_subsequent_mask(tgt.shape[0], tgt.device)
+        outputs = self.decoder(
+            tgt=tgt,
+            tgt_mask=tgt_mask,
+            tgt_key_padding_mask=batch.description_token_ids_mask[:, :-1],
+            memory=encoded_batch.context_encodings,
+            memory_key_padding_mask=encoded_batch.context_encodings_mask,
+        )
+        return self.output(outputs)
+
+    def predict(self, encoded_batch: EncodedBatch, decoding_strategy: str):
+        batch_size = encoded_batch.context_encodings.shape[1]
+        tgt = torch.tensor(
+            [BOS_ID] * batch_size, device=encoded_batch.context_encodings.device
+        ).unsqueeze(dim=0)
+        tgt_mask = Decoder.generate_square_subsequent_mask(self.max_seq_len, tgt.device)
+        pred_all = []
+        for idx in range(self.max_seq_len):
+            tgt_emb = self.pos_encoder(self.embedding(tgt))
+            outputs = self.decoder(
+                tgt_emb,
+                tgt_mask=tgt_mask[: idx + 1, : idx + 1],
+                memory=encoded_batch.context_encodings,
+                memory_key_padding_mask=encoded_batch.context_encodings_mask,
+            )
+            logits = self.output(outputs[-1])
+
+            if decoding_strategy == "greedy":
+                pred_step = logits.argmax(dim=1).tolist()
+            elif decoding_strategy == "nucleus":
+                pred_step = [
+                    helpers.top_k_top_p_sampling(logits[i], top_p=0.95)
+                    for i in range(batch_size)
+                ]
+            else:
+                raise NotImplementedError
+            for b in range(batch_size):
+                if pred_all and pred_all[-1][b].item() in [EOS_ID, PAD_ID]:
+                    pred_step[b] = PAD_ID
+            if all([pred == PAD_ID for pred in pred_step]):
+                break
+            pred_step = torch.tensor(pred_step, device=tgt.device)
+            pred_all.append(pred_step)
+
+            if idx < self.max_seq_len - 1:
+                tgt_step = pred_step.unsqueeze(dim=0)
+                tgt = torch.cat([tgt, tgt_step], dim=0)
+
+        preds = torch.stack(pred_all)
+        return preds
+
+    @staticmethod
+    def generate_square_subsequent_mask(sz: int, device: torch.device) -> torch.Tensor:
+        r"""
+        Generate a square mask for the sequence. The masked positions are filled with
+          float('-inf').
+        Unmasked positions are filled with float(0.0).
+        """
+        return torch.triu(
+            torch.full((sz, sz), float("-inf"), device=device), diagonal=1
+        )

vocab.cond.vocab

@@ -0,0 +1,31 @@
+<pad>	0
+<s>	0
+</s>	0
+<unk>	0
+▁<a>	-1.261
+▁<c>	-1.82603
+<3>	-2.17158
+<0>	-2.26491
+<1>	-2.34
+<2>	-2.36126
+▁<b>	-2.89
+<4>	-3.31157
+<5>	-3.54753
+<6>	-5.02554
+<7>	-5.3972
+<11>	-5.51923
+<8>	-6.03597
+<9>	-6.14342
+<10>	-6.17248
+<13>	-6.28137
+<12>	-6.84099
+<17>	-7.69198
+<14>	-7.72356
+<15>	-8.15065
+<20>	-9.37115
+<18>	-9.52068
+<19>	-9.52068
+<16>	-9.55347
+<23>	-9.58737
+<25>	-9.95894
+<24>	-12.9547

vocab.py

@@ -0,0 +1,41 @@
+import tempfile
+from argparse import ArgumentParser
+
+import sentencepiece as spm
+from transformers.models.bert.tokenization_bert import BertTokenizer
+
+# Load the text tokenizer
+tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
+
+BOS_TOKEN = tokenizer.cls_token
+EOS_TOKEN = tokenizer.sep_token
+UNK_TOKEN = tokenizer.unk_token
+PAD_ID = tokenizer.pad_token_id
+BOS_ID = tokenizer.cls_token_id
+EOS_ID = tokenizer.sep_token_id
+UNK_ID = tokenizer.unk_token_id
+
+# Build the condition (attribute) tokenizer
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    # fmt: off
+    parser.add_argument("--input", nargs="+", required=True)
+    parser.add_argument("--vocab-file", type=str, required=True)
+    parser.add_argument("--vocab-size", type=int, default=31)
+    parser.add_argument("--algo", type=str, default="bpe", choices=["bpe", "word"])
+    # fmt: on
+    args = parser.parse_args()
+    print("Building token vocabulary")
+    with tempfile.NamedTemporaryFile("w") as f:
+        # concatenate input files
+        for input_fname in args.input:
+            with open(input_fname) as input_f:
+                f.write(input_f.read() + "\n")
+        # run sentence piece with bpe
+        spm.SentencePieceTrainer.Train(
+            f"--add_dummy_prefix=false --pad_id=0 --bos_id=1 --eos_id=2 --unk_id=3 "
+            f"--vocab_size={args.vocab_size} "
+            f"--model_prefix={args.vocab_file} --model_type={args.algo} "
+            f"--input={f.name}"
+        )
+~