|
import gradio as gr |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
model_path = "ibm-granite/granite-3.0-1b-a400m-instruct" |
|
tokenizer = AutoTokenizer.from_pretrained(model_path) |
|
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto") |
|
model.eval() |
|
|
|
def generate_response(prompt, max_new_tokens, temperature, top_p, repetition_penalty): |
|
chat = [ |
|
{"role": "user", "content": prompt}, |
|
] |
|
chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True) |
|
|
|
input_tokens = tokenizer(chat, return_tensors="pt").to(model.device) |
|
|
|
output = model.generate( |
|
**input_tokens, |
|
max_new_tokens=max_new_tokens, |
|
temperature=temperature, |
|
top_p=top_p, |
|
repetition_penalty=repetition_penalty, |
|
do_sample=True |
|
) |
|
|
|
response = tokenizer.decode(output[0], skip_special_tokens=True) |
|
return response.split("Human:", 1)[0].strip() |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("# 🙋🏻♂️Welcome to 🌟Tonic's🪨Granite-3.0-1B-A400M-Instruct Demo") |
|
gr.Markdown("Enter a prompt and adjust generation parameters to interact with the 🪨Granite-3.0-1B-A400M-Instruct model.") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=5) |
|
generate_button = gr.Button("Generate Response") |
|
max_new_tokens = gr.Slider(minimum=1, maximum=500, value=100, step=1, label="Max New Tokens") |
|
temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature") |
|
top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top P") |
|
repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, value=1.1, step=0.1, label="Repetition Penalty") |
|
|
|
with gr.Column() : |
|
output = gr.Textbox(label="🪨Granite3-1B", lines=10) |
|
|
|
generate_button.click( |
|
generate_response, |
|
inputs=[prompt, max_new_tokens, temperature, top_p, repetition_penalty], |
|
outputs=output |
|
) |
|
|
|
gr.Markdown("## Examples") |
|
examples = gr.Examples( |
|
examples=[ |
|
["Tell me about the history of artificial intelligence.", 200, 0.7, 0.9, 1.1], |
|
["Write a short story about a robot learning to paint.", 300, 0.8, 0.95, 1.2], |
|
["Explain the concept of quantum computing to a 10-year-old.", 150, 0.6, 0.85, 1.0], |
|
], |
|
inputs=[prompt, max_new_tokens, temperature, top_p, repetition_penalty], |
|
) |
|
|
|
demo.launch() |
