albert-base-v2_mbti-classification / balanced_train_full.py

Upload balanced_train_full.py

beefda1 about 1 year ago

No virus

2.55 kB

	import pandas as pd

	df = pd.read_feather("//media/data/mbti-reddit/disprop_sample100k_total.feather") #change this to proper path
	#'/content/drive/MyDrive/Colab Notebooks/clickbait_hold_X.csv'
	df=df.drop(columns=['authors','subreddit'])

	df=df.sample(80000, random_state=1) #random sampling


	df['labels'] = df['labels'].replace(['INTP','ISTP','ENTP','ESTP','INFP','ISFP','ENFP','ESFP', \
	'INTJ','ISTJ','ENTJ','ESTJ','INFJ','ISFJ','ENFJ','ESFJ'], \
	[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
	df=df.rename(columns={'labels':'labels','comments':'text'})

	from datasets import Dataset

	dataset = Dataset.from_pandas(df)
	dataset.shuffle(seed=27)
	split_set = dataset.train_test_split(test_size=0.2)

	from transformers import AlbertTokenizer, AlbertModel

	tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')

	from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

	model = AutoModelForSequenceClassification.from_pretrained("albert-base-v2", num_labels=16)


	def preprocess_function(examples):
	return tokenizer(examples["text"], truncation=True)

	tokenized_dataset = split_set.map(preprocess_function, batched=True)


	from transformers import DataCollatorWithPadding
	#tokenized_datasets = tokenized_datasets.remove_columns(books_dataset["train"].column_names)
	data_collator = DataCollatorWithPadding(tokenizer=tokenizer)


	import evaluate
	import numpy as np
	def compute_metrics(eval_preds):
	metric = evaluate.combine([

	evaluate.load("precision"),
	evaluate.load("recall")])


	#evaluate.load("precision", average="weighted"),
	#evaluate.load("recall", average="weighted")])

	logits, labels = eval_preds
	predictions = np.argmax(logits, axis=-1)
	return metric.compute(predictions=predictions, references=labels, average='weighted')


	training_args = TrainingArguments(

	evaluation_strategy="epoch",
	#save_strategy="epoch",

	output_dir="/home/deimann/mbti-project/balanced_train",

	#save_total_limit=5,
	#load_best_model_at_end = True,

	learning_rate=2e-5,#2e

	per_device_train_batch_size=36 ,#16

	per_device_eval_batch_size=16,#16

	num_train_epochs=10,

	weight_decay=0.01,

	)

	trainer = Trainer(

	model=model,

	args=training_args,

	train_dataset=tokenized_dataset["train"],

	eval_dataset=tokenized_dataset["test"],

	tokenizer=tokenizer,

	data_collator=data_collator,

	#compute_metrics=compute_metrics,

	)

	trainer.train()