File size: 938 Bytes
b308128
8c245db
 
f9ca505
b308128
7eaa7b0
 
 
 
04fc021
8c245db
 
 
 
7eaa7b0
ac4f141
8c245db
7eaa7b0
 
 
8c245db
 
 
ac4f141
7eaa7b0
5e8be56
8c245db
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
import gradio as gr
import random
import time
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load Vicuna 7B model and tokenizer
model_name = "lmsys/vicuna-7b-v1.3"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

with gr.Blocks() as demo:
    with gr.Row():
        vicuna_chatbot = gr.Chatbot(label="Vicuna", live=True)
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, vicuna_chatbot])

    def respond(message, chat_history, chatbot_idx):
        input_ids = tokenizer.encode(message, return_tensors="pt")
        output = model.generate(input_ids, max_length=50, num_beams=5, no_repeat_ngram_size=2)
        bot_message = tokenizer.decode(output[0], skip_special_tokens=True)
        chat_history.append((message, bot_message))
        time.sleep(2)
        return "", chat_history

    msg.submit(respond, [msg, vicuna_chatbot, 0])

demo.launch()