ReviewAI / backend /data_preprocessing.py
mikahniehaus's picture
Upload backend/data_preprocessing.py with huggingface_hub
15ea0a2 verified
Raw
History Blame Contribute Delete
1.44 kB
import pandas as pd
import json
import os
from sklearn.model_selection import train_test_split
# βœ… Path to Yelp dataset (Update if needed)
data_path = "../yelp_academic_dataset_review.json"
# βœ… List to store processed rows
rows = []
# βœ… Read JSON file (JSON Lines format)
print("πŸ“₯ Loading and processing dataset...")
with open(data_path, "r", encoding="utf-8") as f:
for line in f:
review = json.loads(line)
rows.append({
"text": review["text"],
"stars": int(review["stars"]) - 1, # Convert 1-5 stars to 0-4 for training
"useful": int(round(min(max(review["useful"], 1), 5))), # Normalize usefulness (1-5)
"response": review.get("response", ""), # Keep provided responses, empty if missing
})
# βœ… Convert to Pandas DataFrame
df = pd.DataFrame(rows)
# βœ… Reduce dataset size for quick testing (Use full dataset for actual training)
df = df.sample(100000, random_state=42) # Sample 100K reviews for testing
# βœ… Train/Test Split (80/20 split)
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)
# βœ… Ensure 'data' directory exists
os.makedirs("data", exist_ok=True)
# βœ… Save processed data as CSV
train_df.to_csv("../../data/train.csv", index=False)
test_df.to_csv("../../data/test.csv", index=False)
print("βœ… Data successfully saved to 'data/train.csv' and 'data/test.csv'!")