Johnnyeee/Yelpdata_663

1# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14# TODO: Address all TODOs and remove all explanatory comments15"""TODO: Add a description here."""16 17 18import csv19import json20import os21from typing import List22import datasets23import logging24from random import random25 26# TODO: Add BibTeX citation27# Find for instance the citation on arxiv or on the dataset repo/website28_CITATION = """\29@InProceedings{huggingface:dataset,30title = {A great new dataset},31author={huggingface, Inc.32},33year={2020}34}35"""36 37# TODO: Add description of the dataset here38# You can copy an official description39_DESCRIPTION = """\40This new dataset is designed to solve this great NLP task and is crafted with a lot of care.41"""42 43# TODO: Add a link to an official homepage for the dataset here44_HOMEPAGE = "https://www.yelp.com/dataset/download"45 46# TODO: Add the licence for the dataset here if you can find it47_LICENSE = ""48 49 50class YelpDataset(datasets.GeneratorBasedBuilder):51    """Yelp Dataset focusing on restaurant reviews and business information."""52    53    VERSION = datasets.Version("1.1.0")54    55    BUILDER_CONFIGS = [56        datasets.BuilderConfig(name="restaurants", version=VERSION, description="This part of the dataset covers a wide range of restaurants"),57    ]58 59    DEFAULT_CONFIG_NAME = "restaurants"60    61    _URL = "https://yelpdata.s3.us-west-2.amazonaws.com/"62    _URLS = {63        "business": _URL + "yelp_academic_dataset_business.json",64        "review": _URL + "yelp_academic_dataset_review.json",65    }66 67    def _info(self):68        return datasets.DatasetInfo(69            description=_DESCRIPTION,70            features=datasets.Features({71                "business_id": datasets.Value("string"),72                "name": datasets.Value("string"),73                "address": datasets.Value("string"),74                "city": datasets.Value("string"),75                "state": datasets.Value("string"),76                "postal_code": datasets.Value("string"),77                "latitude": datasets.Value("float"),78                "longitude": datasets.Value("float"),79                "stars_x": datasets.Value("float"),80                "review_count": datasets.Value("float"),81                "is_open": datasets.Value("float"),82                "categories": datasets.Value("string"),83                "hours": datasets.Value("string"),84                "review_id": datasets.Value("string"),85                "user_id": datasets.Value("string"),86                "stars_y": datasets.Value("float"),87                "useful": datasets.Value("float"),88                "funny": datasets.Value("float"),89                "cool": datasets.Value("float"),90                "text": datasets.Value("string"),91                "date": datasets.Value("string"),92                "attributes": datasets.Value("string"),93            }),94            supervised_keys=None,95            homepage="https://www.yelp.com/dataset/download",96            citation=_CITATION,97            license=_LICENSE,98        )99 100    def _split_generators(self, dl_manager: datasets.DownloadManager):101        """Returns SplitGenerators."""102        downloaded_files = dl_manager.download_and_extract(self._URLS)103        104        return [105            datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"files": downloaded_files, "split": "train"}),106            datasets.SplitGenerator(name=datasets.Split.TEST, gen_kwargs={"files": downloaded_files, "split": "test"}),107        ]108 109 110    def _generate_examples(self, files, split):111        """Yields examples as (key, example) tuples."""112        business_path, review_path = files["business"], files["review"]113    114        # Load businesses and filter for restaurants115        with open(business_path, encoding="utf-8") as f:116            businesses = {}117            for line in f:118                business = json.loads(line)119            # Check if 'categories' is not None and contains "Restaurants"120                if business.get("categories") and "Restaurants" in business["categories"]:121                    businesses[business['business_id']] = business122    123    # Generate examples with an attempted 80/20 split for train/test124        with open(review_path, encoding="utf-8") as f:125            for line in f:126                review = json.loads(line)127                business_id = review['business_id']128                if business_id in businesses:129                    business = businesses[business_id]130                    example = {131                        "business_id": business['business_id'],132                        "name": business.get("name", ""),133                        "address": business.get("address", ""),134                        "city": business.get("city", ""),135                        "state": business.get("state", ""),136                        "postal_code": business.get("postal_code", ""),137                        "latitude": business.get("latitude", None),138                        "longitude": business.get("longitude", None),139                        "stars_x": business.get("stars", None),140                        "review_count": business.get("review_count", None),141                        "is_open": business.get("is_open", None),142                        "categories": business.get("categories", ""),143                        "hours": json.dumps(business.get("hours", {})),  # Storing hours as a JSON string144                        "review_id": review.get("review_id", ""),145                        "user_id": review.get("user_id", ""),146                        "stars_y": review.get("stars", None),147                        "useful": review.get("useful", None),148                        "funny": review.get("funny", None),149                        "cool": review.get("cool", None),150                        "text": review.get("text", ""),151                        "date": review.get("date", ""),152                        "attributes": json.dumps(business.get("attributes", {})),  # Storing attributes as a JSON string153                    }154                # Randomly assign to split based on an 80/20 ratio155                    if (split == 'train' and random() < 0.8) or (split == 'test' and random() >= 0.8):156                        yield review['review_id'], example157