# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:torch1.6.0]
#     language: python
#     name: conda-env-torch1.6.0-py
# ---

# +
from json import load
import sys
sys.path.append("/home/ly/workspace/mmsa")

import json
import os
import pickle
import collections
import numpy as np
from typing import *
from tqdm import tqdm
from collections import OrderedDict
from utils.tokenization import BasicTokenizer
from utils.load_yelp import *
seed = 1024
np.random.seed(seed)
# -

base_dir = os.path.join("data","yelp-vistanet")


def load_data():
    with open("data/yelp-vistanet/clear_data.pickle", "rb") as r:
        return pickle.load(r)
class YelpSimpleTokenizer(BasicTokenizer):
    def __init__(self, vocab:Dict[str, int]=None, do_lower_case:bool=True) -> None:
        super(YelpSimpleTokenizer, self).__init__(do_lower_case)
        self.SENT_DELIMITER = '|||'
        self.vocab = vocab
        self.UNK = len(vocab) + 1 if vocab is not None else None # 

    def tokenize(self, text:str) -> List[str]: # 默认切成2d
        res = []
        for sent in text.split(self.SENT_DELIMITER):
            if len(sent) > 0: # 有一定几率出现空字符串
                res.append(super(YelpSimpleTokenizer, self).tokenize(sent))
        return res

    def _getidx(self, token:str):
        return self.vocab.get(token, self.UNK)
        
    def to_idx(self, text:str) -> List[int]:
        assert self.vocab is not None, "No vocab!"
        sents = self.tokenize(text)
        res = []
        for sent in sents:
            res.append([self._getidx(token) for token in sent])
        return res


data = load_data()
vocab = load_glove_vocab()
glove_tokenizer = YelpSimpleTokenizer(vocab["token2idx"], do_lower_case=True)
len(vocab["token2idx"]), len(vocab["idx2token"]), len(vocab["glove_idx"])


# +
def check_photo(i): 
    path = os.path.join(base_dir, "photos", i[:2], i + ".jpg")
    return os.path.exists(path)

def build_glove_data(tokenizer, reviews:List[dict]):
    res = []
    total_img = 0
    for review in tqdm(reviews):
        d = {}
        d["Text"] = tokenizer.to_idx(review["Text"])
        d["Photos"] = []
        for _id in review["Photos"]:
            if check_photo(_id):
                d["Photos"].append(_id)
                total_img += 1
        d["Rating"] = review["Rating"]
        res.append(d)
    print(f"Image num : {total_img}")
    return res


# -

# %%time
glove_data = {}
for key in ["train", "valid", "test"]:
    glove_data[key] = build_glove_data(glove_tokenizer, data[key])

path = os.path.join(base_dir, "glove_data.pickle")
with open(path, "wb") as w:
    pickle.dump(glove_data, w, protocol=pickle.HIGHEST_PROTOCOL)



StarcoderdataJupyterScriptsDedupFiltered