Spaces:
Sleeping
Sleeping
File size: 1,657 Bytes
1623409 |
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 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
import torch
MODEL_NAME = "openbmb/MiniCPM-V-4"
@gr.cache(allow_output_mutation=True)
def load_model():
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True
)
pipeline = TextGenerationPipeline(
model=model,
tokenizer=tokenizer,
device=model.device.index if torch.cuda.is_available() else -1
)
return pipeline
def generate(prompt: str, max_length: int = 100, top_k: int = 50, top_p: float = 0.95):
pipe = load_model()
output = pipe(
prompt,
max_length=max_length,
do_sample=True,
top_k=top_k,
top_p=top_p,
num_return_sequences=1
)
return output[0]["generated_text"]
with gr.Blocks() as demo:
gr.Markdown("# MiniCPM-V-4 Text Generation Demo")
with gr.Row():
prompt_input = gr.Textbox(label="Prompt", placeholder="์ฌ๊ธฐ์ ์
๋ ฅํ์ธ์...", lines=2)
with gr.Row():
max_len = gr.Slider(10, 512, value=100, step=10, label="Max Length")
topk = gr.Slider(1, 100, value=50, step=1, label="Top-k")
topp = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
generate_btn = gr.Button("Generate")
output_box = gr.Textbox(label="Generated Text", lines=5)
generate_btn.click(
fn=generate,
inputs=[prompt_input, max_len, topk, topp],
outputs=output_box
)
demo.launch() |