File size: 3,083 Bytes
93c9911
 
26b527f
 
 
 
 
 
 
 
 
 
 
911bd2d
2d4aaee
93c9911
26b527f
2d4aaee
911bd2d
 
2d4aaee
 
 
 
93c9911
 
911bd2d
 
 
 
26b527f
911bd2d
93c9911
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2d4aaee
 
93c9911
2d4aaee
 
 
 
26b527f
911bd2d
 
93c9911
2d4aaee
93c9911
 
 
2d4aaee
93c9911
 
911bd2d
93c9911
 
 
 
a8537f4
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
"""
Try out gradio.Chatinterface.

colab gradio-chatinterface.

%%writefile reuirements.txt
gradio
transformers
sentencepiece
torch

"""
# pylint: disable=line-too-long, missing-module-docstring, missing-function-docstring
# import torch
import gradio as gr
from examples_list import examples_list
from transformers import AutoModel, AutoTokenizer  # AutoModelForCausalLM,

# device = "cuda" if torch.cuda.is_available() else "cpu"

# tokenizer = AutoTokenizer.from_pretrained("stabilityai/StableBeluga2", use_fast=False)
# model = AutoModelForCausalLM.from_pretrained("stabilityai/StableBeluga2", torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")
# system_prompt = "### System:\nYou are Stable Beluga, an AI that follows instructions extremely well. Help as much as you can. Remember, be safe, and don't do anything illegal.\n\n"
# pipeline = pipeline(task="text-generation", model="meta-llama/Llama-2-7b")

_ = """
tokenizer = AutoTokenizer.from_pretrained(
    "THUDM/chatglm2-6b-int4", trust_remote_code=True
)
chat_model = AutoModel.from_pretrained(
    "THUDM/chatglm2-6b-int4", trust_remote_code=True  # 3.92G
).float()
"""

def stream_chat():
    """samples:

    Sure [('test me', 'Sure')]
    Sure, [('test me', 'Sure,')]
    Sure, I [('test me', 'Sure, I')]
    Sure, I' [('test me', "Sure, I'")]
    Sure, I'd [('test me', "Sure, I'd")]
    """
    resp = ""
    for elm in range(10):
        resp += str(elm)
        from time import sleep
        sleep(0.1)
        yield elm


def chat(message="", history=[]):
    # prompt = f"{system_prompt}### User: {message}\n\n### Assistant:\n"
    # inputs = tokenizer(prompt, return_tensors="pt").to(device=device)
    # output = model.generate(**inputs, do_sample=True, top_p=0.95, top_k=0, max_new_tokens=256)
    # return tokenizer.decode(output[0], skip_special_tokens=True)
    _ = """
    for response, _ in chat_model.stream_chat(
        tokenizer, message, history, max_length=2048, top_p=0.7, temperature=0.95
    ):
       yield response
    """
    g = update_chatbot()
    g.send(None)

    for response in stream_chat():
        # yield response
        g.send(response)

    return history


def update_chatbot():
    while 1:
        message = yield
        print(f"{message=}")


def chat1(message, history):
    # prompt = f"{system_prompt}### User: {message}\n\n### Assistant:\n"
    # inputs = tokenizer(prompt, return_tensors="pt").to(device=device)
    # output = model.generate(**inputs, do_sample=True, top_p=0.95, top_k=0, max_new_tokens=256)
    # return tokenizer.decode(output[0], skip_special_tokens=True)
    for response, _ in chat_model.stream_chat(
        tokenizer, message, history, max_length=2048, top_p=0.7, temperature=0.95
    ):
        yield response, _

with gr.Blocks(theme=gr.themes.Glass(text_size="sm", spacing_size="sm"),) as block:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()

# gr.ChatInterface(
block(
    chat,
    [msg, chatbot],
    [chatbot],
    # title="gradio-chatinterface-tryout",
    # examples=examples_list,
).queue(max_size=2).launch()