| |
| import pandas as pd |
| from tqdm import tqdm |
| from transformers import RobertaForSequenceClassification, RobertaTokenizer |
|
|
| import random |
| random.seed(101) |
|
|
| |
|
|
| 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" |
|
|
|
|
| |
|
|
| 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"{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) |
|
|
| |
|
|
| 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) |
