GPT2 Python Meldungen ausblenden

Question

Ich versuche mit GPT2 Text zu generieren.

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import tensorflow as tf
from transformers import GPT2Tokenizer, TFGPT2LMHeadModel
import sys
 
model_name = "gpt2-medium"
cache_dir = os.getcwd() + "\\gpt2-medium-cache"
tokenizer = GPT2Tokenizer.from_pretrained(model_name, cache_dir=cache_dir)
model = TFGPT2LMHeadModel.from_pretrained(model_name, cache_dir=cache_dir, pad_token_id=tokenizer.eos_token_id)
 
def generate_text(model, tokenizer, start_text, num_tokens):    
    input_ids = tokenizer.encode(start_text, return_tensors='tf')
    output = model.generate(
        input_ids,
        max_length = num_tokens,
        # max_new_tokens = 40,
        num_return_sequences = 1,
        pad_token_id = tokenizer.eos_token_id,
        num_beams = 5,
        no_repeat_ngram_size = 2,
        early_stopping = True,
        top_p = 0.92,
        top_k = 50
    )
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text
 
start_text = sys.argv[1]
num_tokens = int(sys.argv[2])  # Anzahl der gewnschten Tokens im generierten Text
generated_text = generate_text(model, tokenizer, start_text, num_tokens)
 
print(generated_text)

Am Anfang kommt jedes mal so eine Meldung. 

All PyTorch model weights were used when initializing TFGPT2LMHeadModel.  

All the weights of TFGPT2LMHeadModel were initialized from the PyTorch model. If your task is similar to the task the model of the checkpoint was trained on, you can already use TFGPT2LMHeadModel for predictions without further training.

Wie bekomme ich diese Meldung weg?

Answer 1

Du kannst den Transformers Logger auf logging.ERROR stellen, dann werden nur noch Fehler ausgegeben.

import logging
logging.getLogger("transformers").setLevel(logging.ERROR)