|
|
import gradio as gr |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
import torch |
|
|
|
|
|
MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" |
|
|
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) |
|
|
if tokenizer.pad_token is None: |
|
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
MODEL_NAME, |
|
|
torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32 |
|
|
) |
|
|
if torch.cuda.is_available(): |
|
|
model.to("cuda") |
|
|
model.eval() |
|
|
|
|
|
def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50): |
|
|
if not prompt: |
|
|
return "Please enter a prompt." |
|
|
|
|
|
messages = [{"role": "user", "content": prompt}] |
|
|
encoded = tokenizer.apply_chat_template( |
|
|
messages, |
|
|
add_generation_prompt=True, |
|
|
return_tensors="pt", |
|
|
padding=True, |
|
|
return_attention_mask=True, |
|
|
) |
|
|
|
|
|
input_ids = encoded["input_ids"] |
|
|
attention_mask = encoded["attention_mask"] |
|
|
|
|
|
if torch.cuda.is_available(): |
|
|
input_ids = input_ids.to("cuda") |
|
|
attention_mask = attention_mask.to("cuda") |
|
|
|
|
|
output_ids = model.generate( |
|
|
input_ids, |
|
|
attention_mask=attention_mask, |
|
|
max_new_tokens=max_new_tokens, |
|
|
do_sample=True, |
|
|
temperature=temperature, |
|
|
top_k=top_k, |
|
|
pad_token_id=tokenizer.eos_token_id |
|
|
) |
|
|
|
|
|
response = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True) |
|
|
return response |
|
|
|
|
|
|
|
|
demo = gr.Interface( |
|
|
fn=generate_text, |
|
|
inputs=[ |
|
|
gr.Textbox(label="Prompt"), |
|
|
gr.Slider(minimum=10, maximum=500, value=100, label="Max New Tokens"), |
|
|
gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05, label="Temperature"), |
|
|
gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top K") |
|
|
], |
|
|
outputs=gr.Textbox(label="Generated Text"), |
|
|
title="TinyLlama Gradio API", |
|
|
description="Use this via UI or API via `/run/predict`" |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch(server_name="0.0.0.0", server_port=7860) |
|
|
|
