I was wondering how one might go about generating text on a machine with multiple GPU? Here is some example code:
import pandas as pd
from transformers import GPT2Tokenizer, GPT2LMHeadModel
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')
def generate_text(primer_sent: str, chunk_sent_length: int, is_finetuning_current_model=None):
encoded_results = tokenizer(primer_sent, return_tensors='pt')
tmp_ids = encoded_results['input_ids']
max_length = tmp_ids.shape[1] + chunk_sent_length
gen_ids = model.generate(input_ids=tmp_ids,
max_length=max_length,
is_finetuning_current_model=is_finetuning_current_model,
do_sample=True,
top_k=50,
top_p=0.95,
pad_token_id=50256)
decoded_str = tokenizer.decode(gen_ids, skip_special_tokens=True)
return decoded_str
def generate_text_samples(data_filename: str, chunk_sent_length: int):
sent_df = pd.read_csv(data_filename)
sent_df['gen'] = sent_df[0].apply(generate_text, args=(chunk_sent_length,))
return sent_df
text_df = generate_text_samples('/path/to/data.csv', 50)
When I run this code on my machine, it appears that it is all being done on the CPU. Is it possible to distribute this workflow across the multiple gpu I have? Thanks in advance for your help!