Spaces:

helamouri
/

medichat_assignment

Sleeping

App Files Files Community

medichat_assignment / src /training.py

helamouri

update model

eca6215 almost 2 years ago

Raw History Blame Contribute Delete

715 Bytes

from trl import SFTTrainer
from transformers import TrainingArguments

def train_model(model, tokenizer, train_dataset, dataset_text_field, max_seq_length, dataset_num_proc, packing, training_args):
    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=train_dataset,
        dataset_text_field=dataset_text_field,
        max_seq_length=max_seq_length,
        dataset_num_proc=dataset_num_proc,
        packing=packing,
        args=TrainingArguments(**training_args),
    )
    
    # Train the model
    train_results = trainer.train()
    
    # Optionally, you can return more specific training information if necessary
    return train_results

	from trl import SFTTrainer
	from transformers import TrainingArguments

	def train_model(model, tokenizer, train_dataset, dataset_text_field, max_seq_length, dataset_num_proc, packing, training_args):
	trainer = SFTTrainer(
	model=model,
	tokenizer=tokenizer,
	train_dataset=train_dataset,
	dataset_text_field=dataset_text_field,
	max_seq_length=max_seq_length,
	dataset_num_proc=dataset_num_proc,
	packing=packing,
	args=TrainingArguments(**training_args),
	)

	# Train the model
	train_results = trainer.train()

	# Optionally, you can return more specific training information if necessary
	return train_results