| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import gradio as gr | |
| # Load pre-trained GPT-2 model and tokenizer | |
| model_name = "gpt2-large" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| def generate_text(input_text, max_length=16, num_beams=5, do_sample=False, no_repeat_ngram_size=2): | |
| """ | |
| Generate text based on the given input text. | |
| Parameters: | |
| - input_text (str): The input text to start generation from. | |
| - max_length (int): Maximum length of the generated text. | |
| - num_beams (int): Number of beams for beam search. | |
| - do_sample (bool): Whether to use sampling or not. | |
| - no_repeat_ngram_size (int): Size of the n-gram to avoid repetition. | |
| Returns: | |
| - generated_text (str): The generated text. | |
| """ | |
| # Encode the input text and move it to the appropriate device | |
| input_ids = tokenizer(input_text, return_tensors='pt')['input_ids'] | |
| # Generate text using the model | |
| output = model.generate(input_ids, max_length=max_length, num_beams=num_beams, | |
| do_sample=do_sample, no_repeat_ngram_size=no_repeat_ngram_size) | |
| # Decode the generated output | |
| generated_text = tokenizer.decode(output[0]) | |
| return generated_text | |
| # def generate_text_with_nucleus_search(input_text, max_length=16, do_sample=True, top_p=0.9): | |
| # """ | |
| # Generate text with nucleus sampling based on the given input text. | |
| # Parameters: | |
| # - input_text (str): The input text to start generation from. | |
| # - max_length (int): Maximum length of the generated text. | |
| # - do_sample (bool): Whether to use sampling or not. | |
| # - top_p (float): Nucleus sampling parameter. | |
| # Returns: | |
| # - generated_text (str): The generated text. | |
| # """ | |
| # # Encode the input text and move it to the appropriate device | |
| # input_ids = tokenizer(input_text, return_tensors='pt')['input_ids'] | |
| # # Generate text using nucleus sampling | |
| # output = model.generate(input_ids, max_length=max_length, do_sample=do_sample, top_p=top_p) | |
| # # Decode the generated output | |
| # generated_text = tokenizer.decode(output[0]) | |
| # return generated_text | |
| # Create Gradio input interface | |
| input_text_interface = gr.Textbox(lines=5, label="Input Text", placeholder="Enter text for generation...") | |
| # Create Gradio output interface for regular text generation | |
| output_text_interface1 = gr.Textbox(label="Generated Text (Regular)", placeholder="Generated text will appear here...") | |
| # Interface for regular text generation | |
| interface1 = gr.Interface(generate_text, input_text_interface, output_text_interface1, | |
| title="Text Generation with GPT-2", | |
| description="Generate text using the GPT-2 model with regular generation method.", | |
| allow_flagging="never") | |
| # Create Gradio output interface for text generation with nucleus sampling | |
| # output_text_interface2 = gr.Textbox(label="Generated Text (Nucleus Sampling)", placeholder="Generated text will appear here...") | |
| # # Interface for text generation with nucleus sampling | |
| # interface2 = gr.Interface(generate_text_with_nucleus_search, input_text_interface, output_text_interface2, | |
| # title="Text Generation with Nucleus Sampling", | |
| # description="Generate text using nucleus sampling with the GPT-2 model.", | |
| # allow_flagging="never") | |
| # Launch both interfaces | |
| interface1.launch(share=True) | |
| # interface2.launch(share=True) | |