|
|
|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
model_name = "RWKV/rwkv-raven-1b5" |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = AutoModelForCausalLM.from_pretrained(model_name) |
|
|
|
|
|
def generate_text(input_text): |
|
input_ids = tokenizer.encode(input_text, return_tensors="pt") |
|
output = model.generate(input_ids, max_length=100, num_return_sequences=1) |
|
generated_text = tokenizer.decode(output[0], skip_special_tokens=True) |
|
return generated_text |
|
|
|
|
|
def gradio_interface(): |
|
text_input = gr.inputs.Textbox(label="输入文本") |
|
text_output = gr.outputs.Textbox(label="生成文本") |
|
|
|
interface = gr.Interface( |
|
fn=generate_text, |
|
inputs=text_input, |
|
outputs=text_output, |
|
title="文本生成器", |
|
description="输入一段文本，生成相应的文本。", |
|
theme="default" |
|
) |
|
return interface |
|
|
|
interface = gradio_interface() |
|
interface.launch() |
|
|
|
|
|
|
