dataloader.py 551 Bytes
Newer Older
1
2
import pandas as pd

3
class DataLoader:
4
    data_location = 'agent/amazon_data/amazon_reviews_us_Camera_v1_00.tsv'
5
6
7
8
9
10
11
12
13
    reviews = pd.read_csv(data_location, sep='\t', error_bad_lines=False)

    def get_reviews(self, product_id):
        return self.reviews[self.reviews['product_id'] == product_id].reset_index(drop=True)

    def get_product_name(self, product_id):
        return self.get_reviews(product_id)['product_title'][0]

    def get_avg_star_rating(self, product_id):
14
        return float(self.get_reviews(product_id)['star_rating'].mean())