File size: 1,524 Bytes
c53c6b7 |
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 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import argparse
# Load the pre-trained model and tokenizer
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)
# Tokenize the input
input_ids = tokenizer.encode(prompt, return_tensors="pt")
# Generate a response
output = model.generate(input_ids, max_length=100, num_return_sequences=1)
# Decode the output
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
return generated_text
def main():
# Create a Gradio interface
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."
)
# Launch the interface
interface.launch()
if __name__ == '__main__':
main()
|