101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
from typing import Optional
|
|
from tqdm import tqdm
|
|
from datasets import load_dataset
|
|
from transformers import AutoTokenizer
|
|
import re
|
|
|
|
BASE_MODEL = "meta-llama/Meta-Llama-3.1-8B-Instruct"
|
|
MIN_TOKENS = 100
|
|
MAX_TOKENS = 141
|
|
|
|
class Item:
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
|
|
PREFIX = "Price is $"
|
|
|
|
title: str
|
|
price: float
|
|
category: str
|
|
token_count: int = 0
|
|
text: Optional[str]
|
|
details: Optional[str]
|
|
prompt: Optional[str] = None
|
|
include = False
|
|
|
|
def __init__(self, data, price, category):
|
|
self.title = data['title']
|
|
self.price = price
|
|
self.category = category
|
|
self.parse(data)
|
|
|
|
def scrub_details(self):
|
|
details = self.details
|
|
removals = ['"Batteries Included?": "No"', '"Batteries Included?": "Yes"', '"Batteries Required?": "No"', '"Batteries Required?": "Yes"', "By Manufacturer", "Item", "Date First", "Package", ":", "Number of", "Best Sellers", "Number", "Product "]
|
|
for remove in removals:
|
|
details = details.replace(remove, "")
|
|
return details
|
|
|
|
|
|
def parse(self, data):
|
|
self.text = self.title + '\n'
|
|
self.text += '\n'.join(data['description'])+ '\n'
|
|
self.details = data['details']
|
|
if self.details:
|
|
self.text += self.scrub_details() + '\n'
|
|
features = '\n'.join(data['features'])
|
|
if features:
|
|
self.text += '\n' + features
|
|
self.text = re.sub(r'[:\[\]"{}【】\s]+', ' ', self.text).strip()
|
|
self.text = self.text.replace(" ,", ",").replace(",,,",",").replace(",,",",")
|
|
tokens = self.tokenizer.encode(self.text, add_special_tokens=False)
|
|
if len(tokens) > MIN_TOKENS:
|
|
tokens = tokens[:MAX_TOKENS]
|
|
self.text = self.tokenizer.decode(tokens)
|
|
self.make_prompt()
|
|
self.count_tokens()
|
|
self.include = True
|
|
|
|
def question(self):
|
|
prompt = "How much is this?\n"
|
|
prompt += f"{self.text}\n"
|
|
return prompt
|
|
|
|
def messages(self):
|
|
return [
|
|
{"role":"system", "content": "You estimate prices to the nearest dollar"},
|
|
{"role":"user", "content": self.question()},
|
|
{"role":"assistant", "content": f"{self.PREFIX}{str(round(self.price))}.00"}
|
|
]
|
|
|
|
def make_prompt(self):
|
|
prompt = self.tokenizer.apply_chat_template(self.messages(), tokenize=False, add_generation_prompt=False)
|
|
groups = prompt.split('\n\n')
|
|
self.prompt = groups[0]+'\n\n'+'\n\n'.join(groups[2:])
|
|
|
|
def count_tokens(self):
|
|
self.token_count = len(self.tokenizer.encode(self.prompt))
|
|
|
|
def tokens_between(self, low, high):
|
|
return self.token_count >= low and self.token_count < high
|
|
|
|
def test_prompt(self):
|
|
return self.prompt.split(self.PREFIX)[0] + self.PREFIX
|
|
|
|
def read_dataset(name):
|
|
print(f"Loading dataset {name}", flush=True)
|
|
dataset = load_dataset("McAuley-Lab/Amazon-Reviews-2023", f"raw_meta_{name}", split="full", trust_remote_code=True)
|
|
results = []
|
|
for data in dataset:
|
|
try:
|
|
price_str = data['price']
|
|
if price_str:
|
|
price = float(price_str)
|
|
if price >= 0.5 and price <= 999.49:
|
|
item = Item(data, price, name)
|
|
if item.include:
|
|
results.append(item)
|
|
except ValueError:
|
|
pass
|
|
print(f"Completed loading {name} with {len(results):,} datapoints", flush=True)
|
|
del dataset
|
|
return results |