Spaces:
Sleeping
Sleeping
| # Step 2: Import necessary libraries | |
| import gradio as gr | |
| from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM | |
| # Step 3: Load the model and tokenizer | |
| model_name = "unsloth/Llama-3.2-1B" | |
| try: | |
| # Attempt to load the tokenizer and model | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| print(f"Successfully loaded model: {model_name}") | |
| except Exception as e: | |
| # Handle errors and notify the user | |
| print(f"Error loading model or tokenizer: {e}") | |
| tokenizer = None | |
| model = None | |
| # Step 4: Use a pipeline for text generation if model is loaded | |
| if model is not None and tokenizer is not None: | |
| text_gen_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer) | |
| else: | |
| text_gen_pipeline = None | |
| # Step 5: Define the text generation function | |
| def generate_text(prompt, max_length=40, temperature=0.8, top_p=0.9, top_k=40, repetition_penalty=1.5, no_repeat_ngram_size=4): | |
| if text_gen_pipeline is None: | |
| return "Model not loaded. Please check the model name or try a different one." | |
| generated_text = text_gen_pipeline(prompt, | |
| max_length=max_length, | |
| temperature=temperature, | |
| top_p=top_p, | |
| top_k=top_k, | |
| repetition_penalty=repetition_penalty, | |
| no_repeat_ngram_size=no_repeat_ngram_size, | |
| num_return_sequences=1) | |
| return generated_text[0]['generated_text'] | |
| # Step 6: Set up the Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## Text Generation with Llama 3.2 - 1B") | |
| gr.Markdown("For more details, check out this [Google Colab notebook](https://colab.research.google.com/drive/1TCyQNWMQzsjit_z3-0jHCQYfFTpawh8r#scrollTo=5-6MhJj0ZVpk).") | |
| prompt_input = gr.Textbox(label="Input (Prompt)", placeholder="Enter your prompt here...") | |
| max_length_input = gr.Slider(minimum=10, maximum=200, value=40, step=10, label="Maximum Length") | |
| temperature_input = gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.1, label="Temperature (creativity)") | |
| top_p_input = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p (nucleus sampling)") | |
| top_k_input = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-k (sampling diversity)") | |
| repetition_penalty_input = gr.Slider(minimum=1.0, maximum=2.0, value=1.5, step=0.1, label="Repetition Penalty") | |
| no_repeat_ngram_size_input = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="No Repeat N-Gram Size") | |
| output_text = gr.Textbox(label="Generated Text") | |
| generate_button = gr.Button("Generate") | |
| generate_button.click(generate_text, | |
| inputs=[prompt_input, max_length_input, temperature_input, top_p_input, top_k_input, repetition_penalty_input, no_repeat_ngram_size_input], | |
| outputs=output_text) | |
| # Step 7: Launch the app | |
| demo.launch() | |