File size: 4,590 Bytes
bbe7b0a
f3b7005
 
 
d4c9a92
6c69482
0464b4c
4ed0b9b
bbe7b0a
 
 
 
 
 
4ed0b9b
 
bbe7b0a
af7806f
bbe7b0a
 
 
4ed0b9b
 
bbe7b0a
abea35b
4ed0b9b
 
 
 
 
 
 
 
abea35b
4ed0b9b
bbe7b0a
 
4ed0b9b
220ce3a
bbe7b0a
 
4ed0b9b
bbe7b0a
 
 
220ce3a
4ed0b9b
 
 
4fbe483
4ed0b9b
fdf8c66
4ed0b9b
 
 
 
 
 
 
 
fdf8c66
220ce3a
4ed0b9b
8724a4d
220ce3a
 
8724a4d
220ce3a
6e999ef
220ce3a
f3b7005
220ce3a
 
 
 
 
 
6e999ef
4ed0b9b
bbe7b0a
4ed0b9b
220ce3a
4ed0b9b
bbe7b0a
4ed0b9b
bbe7b0a
 
 
 
4ed0b9b
bbe7b0a
 
 
4ed0b9b
bbe7b0a
 
4ed0b9b
bbe7b0a
 
 
 
 
 
4ed0b9b
bbe7b0a
 
 
 
 
 
4ed0b9b
 
bbe7b0a
 
 
 
 
 
4ed0b9b
bbe7b0a
0b7e553
 
 
 
4ed0b9b
 
7161b69
b1d449b
bbe7b0a
 
 
 
 
 
 
0b7e553
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import os
from threading import Thread
from typing import Iterator

import gradio as gr
import spaces
import torch
from transformers import pipeline, AutoTokenizer

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

DESCRIPTION = """\
# ZhongJing 2 1.8B Merge
This Space demonstrates model [CMLL/ZhongJing-2-1_8b-merge](https://huggingface.co/CMLL/ZhongJing-2-1_8b-merge) for text generation. Feel free to play with it, or duplicate to run generations without a queue! If you want to run your own service, you can also [deploy the model on Inference Endpoints](https://huggingface.co/inference-endpoints).
"""

LICENSE = """
<p/>
---
As a derivative work of [CMLL/ZhongJing-2-1_8b-merge](https://huggingface.co/CMLL/ZhongJing-2-1_8b-merge),
this demo is governed by the original [license](https://huggingface.co/CMLL/ZhongJing-2-1_8b-merge/LICENSE).
"""

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"

if torch.cuda.is_available():
    model_id = "CMLL/ZhongJing-2-1_8b-merge"
    pipe = pipeline("text-generation", model=model_id)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.use_default_system_prompt = False

@spaces.GPU
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str = "You are a helpful TCM medical assistant named 仲景中医大语言模型, created by 医哲未来.",
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    conversation = [{"role": "system", "content": system_prompt}]
    for user, assistant in chat_history:
        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
    conversation.append({"role": "user", "content": message})

    input_text = "\n".join([f"{entry['role']}: {entry['content']}" for entry in conversation])

    generate_kwargs = {
        "max_new_tokens": max_new_tokens,
        "do_sample": True,
        "top_p": top_p,
        "top_k": top_k,
        "temperature": temperature,
        "repetition_penalty": repetition_penalty,
    }

    # Function to run the generation
    def run_generation():
        try:
            results = pipe(input_text, **generate_kwargs)
            return results
        except Exception as e:
            return [f"Error in generation: {e}"]

    # Run generation in a separate thread and wait for it to finish
    outputs = []
    generation_thread = Thread(target=lambda: outputs.extend(run_generation()))
    generation_thread.start()
    generation_thread.join()

    for output in outputs:
        yield output['generated_text'] if isinstance(output, dict) else output

chat_interface = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.Textbox(label="System prompt", lines=6, value="You are a helpful TCM medical assistant named 仲景中医大语言模型, created by 医哲未来."),
        gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="Temperature",
            minimum=0.1,
            maximum=4.0,
            step=0.1,
            value=0.6,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.2,
        ),
    ],
    stop_btn=None,
    examples=[
        ["Hello there! How are you doing?"],
        ["Can you explain briefly to me what is the Python programming language?"],
        ["Explain the plot of Cinderella in a sentence."],
        ["How many hours does it take a man to eat a Helicopter?"],
        ["Write a 100-word article on 'Benefits of Open-Source in AI research'"],
    ],
)

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
    chat_interface.render()
    gr.Markdown(LICENSE)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()