File size: 1,959 Bytes
1b0843f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
from tqdm.auto import tqdm
# Constants
batch_size = 1000
# Load tokenizer and model
model_checkpoint = "PleIAs/French-TV-Headline-Classification"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, model_max_length=512)
model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
classification_pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
# Read the dataset
val_classification = pd.read_parquet("[file]")
val_classification.reset_index(drop=True, inplace=True)
# Calculate the number of batches needed
num_batches = (len(val_classification) + batch_size - 1) // batch_size
# Initialize the list to collect DataFrames
list_df = []
for i in tqdm(range(num_batches), desc="Processing batches"):
start_index = i * batch_size
end_index = min((i + 1) * batch_size, len(val_classification))
batch = val_classification.iloc[start_index:end_index]
# Extract texts from the DataFrame
texts = batch["corrected_text"].tolist()
# Classify texts in batches
classifications = classification_pipeline(texts, truncation=True, padding=True, top_k=None)
# Prepare data for DataFrame
rows = []
for text_index, class_results in enumerate(classifications):
for entry in class_results:
rows.append({
'text_id': start_index + text_index,
'label': entry['label'],
'score': round(entry['score'] * 100, 2),
'identifier': batch.iloc[text_index]['identifier']
})
# Create DataFrame from the processed batch
df = pd.DataFrame(rows)
list_df.append(df)
# Concatenate all DataFrames in the list
final_df = pd.concat(list_df, ignore_index=True)
print(final_df)
# Save the resulting DataFrame to a CSV file
final_df.to_csv("transcript_classification.csv", index=False)
|