Create example_classify_script_for_csv.py

example_classify_script_for_csv.py

@@ -0,0 +1,55 @@
+# CUDA_VISIBLE_DEVICES=0,1,2,3 python classify_generics.py
+import pandas as pd
+from tqdm import tqdm
+from transformers import RobertaForSequenceClassification, RobertaTokenizer
+
+import random
+random.seed(101)
+
+## PARAMETERS ##
+
+input_csv = "data/input.csv"
+text_column = "text_column_name"
+output_csv = ""
+output_csv = f"{input_csv}_scored_for_generics.csv" if not output_csv else output_csv
+
+
+model_path = "ilyocoris/generics-classifier-mgen"
+
+
+## LOAD DATA & MODEL ##
+
+data = pd.read_csv(input_csv)
+data = data.to_dict(orient="records")
+
+def load_roberta_classifier(model_name, checkpoint=None, base_model="roberta-large"):
+    # model_path = f"data/training_runs/{model_name}/models/checkpoint-{checkpoint}" if checkpoint else f"data/models/{model_name}"
+    model_path = f"{model_name}/checkpoint-{checkpoint}" if checkpoint else model_name
+    model = RobertaForSequenceClassification.from_pretrained(
+        model_path,
+        device_map="auto",
+        num_labels=1
+    )
+    tokenizer = RobertaTokenizer.from_pretrained(base_model)
+    return model, tokenizer
+
+def classify_batch(batch_sentences, model, tokenizer):
+    inputs = tokenizer(batch_sentences, return_tensors="pt", padding=True, truncation=True)
+    inputs = {k: v.cuda() for k, v in inputs.items()}
+    outputs = model(**inputs)
+    return outputs.logits.squeeze().cpu().detach().numpy()
+
+model, tokenizer = load_roberta_classifier(model_path) 
+
+## RUN CLASSIFIER ##
+
+scores = []
+batch_size = 16
+for i in tqdm(range(0, len(data), batch_size)):
+    batch = data[i:i+batch_size]
+    sentences = [d[text_column]  for d in batch]
+    scores.extend(classify_batch(sentences, model, tokenizer).tolist())
+
+df = pd.DataFrame(data)
+df["score"] = scores
+df.to_csv(f"{output_csv}", index=False)