|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
import argparse |
|
|
|
|
|
models = { |
|
"EleutherAI/gpt-neo-2.7B": "EleutherAI/gpt-neo-2.7B", |
|
"BlinkDL/rwkv-4-pile-430m": "BlinkDL/rwkv-4-pile-430m", |
|
"BlinkDL/rwkv-4-pile-1b5": "BlinkDL/rwkv-4-pile-1b5", |
|
"BlinkDL/RWKV-4-Raven": "BlinkDL/RWKV-4-Raven" |
|
} |
|
|
|
def generate_text(prompt, model_name): |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = AutoModelForCausalLM.from_pretrained(model_name) |
|
|
|
|
|
input_ids = tokenizer.encode(prompt, 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 main(): |
|
|
|
model_dropdown = gr.inputs.Dropdown(choices=list(models.keys()), label="Select Model") |
|
prompt_input = gr.inputs.Textbox(lines=5, placeholder="Enter your prompt here...", label="Prompt") |
|
output_text = gr.outputs.Textbox(label="Generated Text") |
|
|
|
interface = gr.Interface( |
|
fn=generate_text, |
|
inputs=[prompt_input, model_dropdown], |
|
outputs=output_text, |
|
title="Chat-bot using RWKV LLM", |
|
description="Select a model and enter a prompt to generate text using the chat-bot." |
|
) |
|
|
|
|
|
interface.launch() |
|
|
|
if __name__ == '__main__': |
|
main() |
|
|
