sawaw/scripts/experiment.py

35 lines
955 B
Python

from dataclasses import dataclass
from random import sample
from typing import List
import pandas as pd
from sawaw import SAWAWEntry, analyze_gpt3_5
# Create a DataFrame from ./combined_data_with_aspects.csv
df = pd.read_csv("data/combined_data_with_aspects.csv")
# Keep only text and aspect columns
df = df[["text", "aspect"]]
# Remove rows with NaN values
df = df.dropna()
# Remove rows with empty strings
df = df[df["text"] != ""]
# Create a list of SAWAWEntry objects
entries = []
for index, row in df.iterrows():
comment = row["text"]
aspect_words = row["aspect"].split(",")
aspect_words = [aspect_word.strip() for aspect_word in aspect_words]
entries.append(SAWAWEntry(comment, aspect_words))
# Pick 10 random entries
entries = sample(entries, 3)
# Query the entries
for entry in entries:
print(entry)
analyze_gpt3_5(entry)
SAWAWEntry.print_legends()
print(entry)
input("Press Enter to continue...")