Package: textrl
TextRL: Text Generation with Reinforcement LearningTextRL is a Python library that aims to improve text generation using reinforcement learning, building upon Hugging Face's Transformers, PFRL, and OpenAI GYM. TextRL is designed to be easily customizable and can be applied to various text-generation models.Table of ContentsIntroductionExamplesGPT-2 ExampleFLAN-T5 ExampleBigscience/BLOOMZ-7B1-MT Example176B BLOOM ExampleControllable Generation via RL ExampleInstallationPip InstallBuild from SourceUsageInitialize Agent and EnvironmentSetup Reward Function for EnvironmentPrepare for TrainingTrainingDump ModelKey Parameters for RL TrainingIntroductionTextRL utilizes reinforcement learning to fine-tune text generation models. It is built upon the following libraries:Hugging Face's TransformersPFRLOpenAI GYMExample -gpt2CLICK MEGPT2 ExampleimportpfrlfromtextrlimportTextRLEnv,TextRLActor,train_agent_with_evaluationfromtransformersimportAutoModelForCausalLM,AutoTokenizerimportloggingimportsyslogging.basicConfig(level=logging.INFO,stream=sys.stdout,format='')checkpoint="gpt2"tokenizer=AutoTokenizer.from_pretrained(checkpoint)model=AutoModelForCausalLM.from_pretrained(checkpoint,torch_dtype="auto",device_map="auto")model=model.cuda()classMyRLEnv(TextRLEnv):defget_reward(self,input_item,predicted_list,finish):# predicted will be the list of predicted tokenreward=[0]iffinish:reward=[1]# calculate reward score base on predicted_listreturnrewardobservaton_list=[{"input":"explain how attention work in seq2seq model"}]env=TextRLEnv(model,tokenizer,observation_input=observaton_list,max_length=20,compare_sample=2)actor=TextRLActor(env,model,tokenizer,act_deterministically=False,temperature=1.0,top_k=0,top_p=1.0,repetition_penalty=2)agent=actor.agent_ppo(update_interval=2,minibatch_size=2,epochs=10)print(actor.predict(observaton_list[0]))train_agent_with_evaluation(agent,env,steps=100,eval_n_steps=None,eval_n_episodes=1,eval_interval=2,outdir='bloom—test',)print(actor.predict(observaton_list[0]))Example -flan-t5CLICK MEExample Codecolab
example:google/flan-t5-baseimportpfrlfromtextrlimportTextRLEnv,TextRLActor,train_agent_with_evaluationfromtransformersimportAutoModelForSeq2SeqLM,AutoTokenizerimportloggingimportsyslogging.basicConfig(level=logging.INFO,stream=sys.stdout,format='')tokenizer=AutoTokenizer.from_pretrained("google/flan-t5-base")model=AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")model.eval()model.cuda()sentiment=pipeline('sentiment-analysis',model="cardiffnlp/twitter-roberta-base-sentiment",tokenizer="cardiffnlp/twitter-roberta-base-sentiment",device=0,return_all_scores=True)classMyRLEnv(TextRLEnv):defget_reward(self,input_item,predicted_list,finish):# predicted will be the list of predicted tokenreward=0iffinishorlen(predicted_list[0])>=self.env_max_length:predicted_text=tokenizer.convert_tokens_to_string(predicted_list[0])# sentiment classifierreward=sentiment(input_item['input']+predicted_text)[0][0]['score']*10returnrewardobservaton_list=[{'input':'i think dogecoin is'}]env=MyRLEnv(model,tokenizer,observation_input=observaton_list,compare_sample=1)actor=TextRLActor(env,model,tokenizer,optimizer='adamw',temperature=0.8,top_k=100,top_p=0.85,)agent=actor.agent_ppo(update_interval=50,minibatch_size=3,epochs=10,lr=3e-4)print(actor.predict(observaton_list[0]))pfrl.experiments.train_agent_with_evaluation(agent,env,steps=3000,eval_n_steps=None,eval_n_episodes=1,train_max_episode_len=100,eval_interval=10,outdir='checkpoint',)agent.load("./checkpoint/best")print(actor.predict(observaton_list[0]))Example -bigscience/bloomz-7b1-mtCLICK MEbloomz-7b1-mt ExampleimportpfrlfromtextrlimportTextRLEnv,TextRLActor,train_agent_with_evaluationfromtransformersimportAutoModelForCausalLM,AutoTokenizerimportloggingimportsyslogging.basicConfig(level=logging.INFO,stream=sys.stdout,format='')checkpoint="bigscience/bloomz-7b1-mt"tokenizer=AutoTokenizer.from_pretrained(checkpoint)model=AutoModelForCausalLM.from_pretrained(checkpoint,torch_dtype="auto",device_map="auto")model=model.cuda()classMyRLEnv(TextRLEnv):defget_reward(self,input_item,predicted_list,finish):# predicted will be the list of predicted tokenreward=[0]iffinish:reward=[1]# calculate reward score base on predicted_listreturnrewardobservaton_list=[{"input":"explain how attention work in seq2seq model"}]env=TextRLEnv(model,tokenizer,observation_input=observaton_list,max_length=20,compare_sample=2)actor=TextRLActor(env,model,tokenizer,act_deterministically=False,temperature=1.0,top_k=0,top_p=1.0)agent=actor.agent_ppo(update_interval=2,minibatch_size=2,epochs=10)print(actor.predict(observaton_list[0]))train_agent_with_evaluation(agent,env,steps=100,eval_n_steps=None,eval_n_episodes=1,eval_interval=2,outdir='bloom—test',)print(actor.predict(observaton_list[0]))Example - 176B BLOOMCLICK MEbloomz-176B ExampleStrongly recommend contribute on public swarm to increase petals capacityhttps://github.com/bigscience-workshop/petalsinstallpip install petals -UfirstimportpfrlfromtextrlimportTextRLEnv,TextRLActor,train_agent_with_evaluationfromtransformersimportBloomTokenizerFastfrompetalsimportDistributedBloomForCausalLMimportloggingimportsyslogging.basicConfig(level=logging.INFO,stream=sys.stdout,format='')MODEL_NAME="bigscience/bloom-petals"tokenizer=BloomTokenizerFast.from_pretrained(MODEL_NAME)model=DistributedBloomForCausalLM.from_pretrained(MODEL_NAME)model=model.cuda()classMyRLEnv(TextRLEnv):defget_reward(self,input_item,predicted_list,finish):# predicted will be the list of predicted tokenreward=[0]iffinish:reward=[1]# calculate reward score base on predicted_listreturnrewardobservaton_list=[{"input":"explain how attention work in seq2seq model"}]env=TextRLEnv(model,tokenizer,observation_input=observaton_list,max_length=20,compare_sample=2)actor=TextRLActor(env,model,tokenizer,act_deterministically=False,temperature=1.0,top_k=0,top_p=1.0)agent=actor.agent_ppo(update_interval=2,minibatch_size=2,epochs=10)print(actor.predict(observaton_list[0]))train_agent_with_evaluation(agent,env,steps=100,eval_n_steps=None,eval_n_episodes=1,eval_interval=2,outdir='bloom—test',)print(actor.predict(observaton_list[0]))Example - Controllable generation via RL to let Elon Musk speak ill of DOGECLICK ME[Controllable generation via RL to let Elon Musk speak ill of DOGE
](https://github.com/voidful/TextRL/blob/main/example/2022-12-10-textrl-elon-musk.ipynb)colab
example:bigscience/bloom-560mcolab
exmaple:huggingtweets/elonmuskbefore:i think dogecoin is a great idea.after:i think dogecoin is a great idea, but I think it is a little overused.Installationpip installpipinstallpfrl@git+https://github.com/voidful/pfrl.git
pipinstalltextrlBuild from sourcegit clone and cd into this project.pipinstall-e.UsageInitialize agent and environmentimporttorchfromtextrlimportTextRLEnv,TextRLActor,train_agent_with_evaluationfromtransformersimportAutoModelForCausalLM,AutoTokenizercheckpoint="bigscience/bloomz-7b1-mt"tokenizer=AutoTokenizer.from_pretrained(checkpoint)model=AutoModelForCausalLM.from_pretrained(checkpoint,torch_dtype="auto",device_map="auto")model=model.cuda()Set up reward function for environmentpredicted(list[str]): will be the list of predicted tokensfinish(bool): whether the end of sentence has been reached or notclassMyRLEnv(TextRLEnv):defget_reward(self,input_item,predicted_list,finish):iffinish:reward=[0]# calculate reward score based on predicted_listreturnrewardPrepare for trainingobservation_list should be a list of all possible input strings for model trainingExample:observation_list = [{"input":'testing sent 1'},{"input":'testing sent 2'}]env=MyRLEnv(model,tokenizer,observation_input=observation_list)actor=TextRLActor(env,model,tokenizer)agent=actor.agent_ppo(update_interval=10,minibatch_size=2000,epochs=20)Trainn_episodes=1000max_episode_len=200# max sentence lengthforiinrange(1,n_episodes+1):obs=env.reset()R=0t=0whileTrue:action=agent.act(obs)obs,reward,done,pred=env.step(action)R+=rewardt+=1reset=t==max_episode_lenagent.observe(obs,reward,done,reset)ifdoneorreset:breakifi%10==0:print('episode:',i,'R:',R)ifi%50==0:print('statistics:',agent.get_statistics())print('Finished.')Another way to train:importloggingimportsyslogging.basicConfig(level=logging.INFO,stream=sys.stdout,format='')train_agent_with_evaluation(agent,env,steps=1000,eval_n_steps=None,eval_n_episodes=1500,train_max_episode_len=50,eval_interval=10000,outdir='somewhere',)Predictionagent.load("somewhere/best")# loading the best modelactor.predict("input text")This updated usage section provides a comprehensive guide on how to initialize the agent and environment, set up the reward function for the environment, prepare for training, train the model, and make predictions. It also includes an alternative way to train the model using thetrain_agent_with_evaluationfunction.Dump trained model to huggingface's modeltextrl-dump--model./model_path_before_rl--rl./rl_path--dump./output_dirKey Parameters for RL TrainingTo finetune a language model using RL, you need to modify the reward function:fromtextrlimportTextRLEnvclassMyRLEnv(TextRLEnv):defget_reward(self,input_item,predicted_list,finish):# input_item is the prompt input for the model, it will be one of your observation# an observation will be a list of sentence of eg: ['inputted sentence','xxx','yyy']# only the first input will feed to the model 'inputted sentence', and# the remaining can be the reference for reward calculation# predicted_list is the list of predicted sentences of RL model generated,# it will be used for ranking reward calculation# finish is the end of sentences flags, get_reward will be called during generating each word, and# when finish is True, it means the sentence is finished, it will use for sentence level reward calculation.# reward should be the list equal to the length of predicted_listreturnrewardParameters for sampling diverse examples:actor=TextRLActor(env,model,tokenizer,act_deterministically=False,# select the max probability token for each step or nottemperature=1,# temperature for samplingcompare_sample=2,# num of sample to ranktop_k=0,# top k samplingtop_p=1.0,)# top p samplingWhen training a reinforcement learning (RL) model, several key parameters need to be tuned to ensure optimal performance. Here is a list of important parameters and their descriptions:Update Interval: This determines how often the RL agent updates its policy based on collected experiences. A smaller update interval means the agent learns more frequently from recent experiences, while a larger interval allows more experiences to accumulate before learning. In the example above, the update interval is set to 10.update_interval=10Minibatch Size: The number of experiences sampled from the experience replay buffer to compute the gradient update. A larger minibatch size helps to stabilize learning and reduce variance, but at the cost of increased computational requirements.minibatch_size=2000Epochs: The number of times the agent iterates through the entire minibatch to update its policy. More epochs can lead to better learning but may increase the risk of overfitting.epochs=20Discount Factor (Gamma): This parameter determines how much future rewards are discounted when calculating the expected return. A value closer to 1 makes the agent more farsighted, while a value closer to 0 makes the agent more focused on immediate rewards.gamma=0.99Learning Rate: The step size used for updating the policy. A larger learning rate allows for faster convergence but may lead to instability in learning, while a smaller learning rate ensures stable learning at the cost of slower convergence.lr=1e-4Epsilon: A parameter used in the PPO algorithm to clip the policy ratio. This helps to control the magnitude of policy updates, preventing excessively large updates that can destabilize learning.epsilon=0.2Entropy Coefficient: This parameter encourages exploration by adding a bonus reward for taking less certain actions. A higher entropy coefficient promotes more exploration, while a lower coefficient focuses the agent on exploiting known strategies.entropy_coef=0.01Training Steps: The total number of steps the agent takes during training. More steps typically lead to better learning but may require more computational time.steps=1000Evaluation Interval: The number of training steps between evaluations. Increasing the evaluation interval reduces the computational time spent on evaluation, but it may also reduce the frequency at which you can monitor the agent's progress.eval_interval=10000Max Episode Length: The maximum number of steps allowed in a single episode during training. This can prevent the agent from getting stuck in long, unproductive episodes.train_max_episode_len=50These parameters need to be carefully tuned based on the specific problem and environment to achieve the best performance. It is generally recommended to start with default values and then adjust them based on the observed learning behavior.