<jupyter_start><jupyter_text>Dataset Builder> Builds a number of types of datasets that are augmented by the `wikidata` knowledge base Dataset VariationsKeywords can be supplied to the `build` function through the `dataset_type` keyword argument Description`DatasetBuilder.build(ds, dataset_type='description')`Augements the given dataset with the description of the `keyword`. *Example*`<NAME> is my favorite basketball player. {<NAME>: {Description: American basketball player}}`<jupyter_code>#export
class DatasetBuilder():
    "Build a dataset using `get_entities_in_text`"
    def __init__(self):
#         self.db = WikiDatabase()
        self.nlp = spacy.load("en_core_web_sm", disable=["tok2vec", "tagger", "parser", "attribute_ruler", "lemmatizer"])

#         module_url = "https://tfhub.dev/google/universal-sentence-encoder/4" #@param ["https://tfhub.dev/google/universal-sentence-encoder/4", "https://tfhub.dev/google/universal-sentence-encoder-large/5"]
#         self.encoder = hub.load(module_url)
    
    def build(self, ds, dataset_type='random'):
        "Build a database based a given dataset"
        if dataset_type == 'random':
            return ds.map(self.random, batched=False)
        elif dataset_type == 'description':
            return ds.map(self.description, batched=False)
        elif dataset_type == 'relevant':
            pass
        
    def build_knowledge_entities(self, ds, split):
        ds = ds.map(self.get_entities, batched=True, num_proc=4)
        ds.save_to_disk('data/augmented_datasets/entities/' + split + '/')
        
    def build_csv(self, ds, split):
        ds = ds.map(self.retrieve_knowledge, batched=False)
        ds.save_to_disk('data/augmented_datasets/')
        
    def get_entities(self, batch):
        import pdb; pdb.set_trace()
        doc = self.nlp(batch)
        entities = doc.ents
        
    def retrieve_knowledge(self, sequence):
        text = sequence['text']
        entities = self.get_entities_in_text(text)
        knowledge = self.add_associations(entities)
        sequence['knowledge'] = knowledge
        return sequence
            
    def add_associations(self, entities):
        "Returns list of entity/association dictionaries"
        associations = []
        for e in entities:
            a = self.get_entity_associations(e)
            k = {e[1]: a}
            associations.append(k)
        return associations
        
    def _get_json(self, item):
        """Return JSON version of list object"""
        d = {"label": None, "description": None}
        d['label'] = item[1]
        d['description'] = item[2]
        return json.dumps(d)
    
    def get_entities_in_text(self, text):
        "Returns entities found in the sentence `text`"
        doc = self.nlp(text)
        entities = []
        spacy_entities = doc.ents
        for entity in spacy_entities: 
            entity = self.db.get_entity_by_label(entity.text)
            if entity:
                entities.append(entity)
        return entities  
    
    def get_entity_associations(self, entity):
        """
        Given an `entity_id` return a dictionary containing all the associated properties.
        """
        entity_id = entity[0]
        entity_associations_dict = {'id':entity_id, 'description':entity[2]}
        # Remove all None values from list
        associations = self.db.get_entity_associations(entity_id)
        if not associations:
            return None
        for property_id, related_entity_id in associations: 
            property_name, related_entity_label = self.db.get_property_string(property_id, related_entity_id)
            entity_associations_dict[property_name] = related_entity_label
        return entity_associations_dict<jupyter_output><empty_output><jupyter_text>Testing<jupyter_code># Build description dataset
from kirby.run_params import RunParams
from kirby.data_manager import DataManager
from datasets import load_dataset

run_params = RunParams(debug=True)
data_manager = DataManager(run_params)
ds_builder = DatasetBuilder()

split = 'train'
ds = data_manager.load(split)
ds_builder.build_knowledge_entities(ds, split)

split = 'valid'
ds = data_manager.load(split)
ds_builder.build_knowledge_entities(ds, split)
ds
# creation
ds_builder = DatasetBuilder()

assert isinstance(ds_builder, DatasetBuilder)
# Test ranked phrases
x = "<NAME> is my favorite basketball player."
ds_builder.rake.extract_keywords_from_text(x)
ranked_phrases = ds_builder.rake.get_ranked_phrases()
assert ranked_phrases == ['favorite basketball player', 'stephen curry'], "RAKE failed"
print(ds_builder.db.get_entity_by_label('<NAME>'))
assert ds_builder.db.get_entity_by_label('<NAME>') == ['Q11571', '<NAME>', 'Portuguese association football player'], 'ERROR in `database_proxy`'
# Get Entities from the sentence
x = "Microsoft has bought Bethesda"
entities = ds_builder.get_entities_in_text(x)
print(entities)
assert entities == [['Q2283', 'Microsoft', 'American multinational technology corporation'],\
                    ['Q224892', 'Bethesda', 'Wikimedia disambiguation page']],\
                    'Error in `dataset_builder.get_entities_in_text()`'
# Get associations from an entity
associations = ds_builder.get_entity_associations(entities[0][0])
assert associations == {"topic's main Wikimedia portal": 'Portal:Microsoft',
 'founded by': '<NAME>',
 'country': 'United States of America',
 'instance of': 'software company',
 'headquarters location': 'Redmond',
 'stock exchange': 'NASDAQ',
 'chief executive officer': '<NAME>',
 "topic's main category": 'Category:Microsoft',
 'subsidiary': 'Xbox Game Studios',
 'described by source': 'Lentapedia (full versions)',
 'industry': 'technology industry',
 'product or material produced': 'Microsoft Windows',
 'legal form': 'Washington corporation',
 'business division': 'Microsoft Research',
 'history of topic': 'history of Microsoft',
 'member of': 'Alliance for Open Media',
 'permanent duplicated item': None,
 'part of': 'NASDAQ-100',
 'award received': 'Big Brother Awards',
 'owner of': 'Microsoft TechNet',
 'owned by': 'BlackRock',
 'board member': '<NAME>',
 'chairperson': '<NAME>',
 'location of formation': 'Albuquerque',
 'director / manager': '<NAME>',
 'external auditor': 'Deloitte & Touche LLP',
 'partnership with': 'ID2020'}, 'Error in `dataset_builder.get_entity_associations()`'
# Description
text = "Darth Vader cut off <NAME>'s hand"
ds_builder.description(text)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11643308178037887, 'fraction_numerical': 0.005236408439858309, 'mean_word_length': 3.81037037037037, 'pattern_counts': {'":': 4, '<': 21, '<?xml version=': 0, '>': 22, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '290973', 'license': 'Apache-2.0', 'n_tokens_mistral': 1923, 'n_tokens_neox': 1889, 'n_words': 491, 'path': 'notebooks/01_dataset_builder.ipynb', 'repo_name': 'rmorain/kirby', 'size': '6589'}