import pandas as pd
from transformers import pipeline
import sys

# Load your data
df = pd.read_csv(sys.argv[1])  # change sep if needed
df.columns = df.columns.str.strip()  # clean column names

# Check your column names if still unsure
print("Columns:", df.columns.tolist())

# Load NER pipeline with correct aggregation strategy
ner_pipeline = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")

# Define entity extraction function
def extract_entities(text):
    try:
        ner_results = ner_pipeline(str(text))
        return [(ent['word'], ent['entity_group']) for ent in ner_results]
    except Exception as e:
        return []

# Apply to 'comment' column
df['entities'] = df['comment'].apply(extract_entities)

# Save result
df.to_csv(sys.argv[1], index=False)