File size: 1,760 Bytes
69cea9f 1f82f00 2a8bae1 1f82f00 2a8bae1 1f82f00 2a8bae1 69cea9f 1f82f00 69cea9f 2a8bae1 69cea9f 1f82f00 69cea9f 1f82f00 69cea9f 1f82f00 |
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 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "fla-hub/rwkv7-2.9B-world"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.float32,
low_cpu_mem_usage=True,
device_map="cpu"
)
print("Model loaded!")
def respond(message, history, system_message, max_tokens, temperature, top_p):
messages = [{"role": "system", "content": system_message}]
for human, assistant in history:
messages.append({"role": "user", "content": human})
messages.append({"role": "assistant", "content": assistant})
messages.append({"role": "user", "content": message})
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
return response
chatbot = gr.ChatInterface(
respond,
type="messages",
additional_inputs=[
gr.Textbox(value="You are a helpful assistant.", label="System message"),
gr.Slider(1, 512, 256, step=1, label="Max tokens"),
gr.Slider(0.1, 2.0, 0.7, step=0.1, label="Temperature"),
gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top-p"),
],
)
chatbot.launch()
|