File size: 2,620 Bytes
4e567e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load model and tokenizer
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() |