File size: 13,978 Bytes
9255bb8
 
 
b1c5198
a09b62a
ae5bb78
9255bb8
 
 
6e12956
 
9255bb8
 
 
ae5bb78
 
b1f91f1
6e12956
 
c6b78e0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6e12956
 
905d227
816cea8
905d227
 
 
 
816cea8
905d227
 
 
 
 
816cea8
905d227
 
 
 
6e12956
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a47d2f3
9255bb8
ca38a58
 
 
 
 
 
6150c8f
a47d2f3
6150c8f
 
 
a47d2f3
6150c8f
 
a47d2f3
cd7fced
4d0f760
cd7fced
4d0f760
a47d2f3
6150c8f
4d0f760
a47d2f3
cd7fced
 
 
a47d2f3
6150c8f
 
a47d2f3
cd7fced
 
 
a47d2f3
6150c8f
4d0f760
b033af5
 
 
a47d2f3
6150c8f
 
cd7fced
 
 
 
 
9255bb8
 
 
 
 
6150c8f
 
 
 
 
 
9255bb8
 
6e12956
905d227
6e12956
 
6150c8f
619704f
6150c8f
6e12956
6150c8f
9255bb8
 
 
 
 
6150c8f
9255bb8
 
 
 
 
 
 
 
c8f6eb0
6e12956
 
 
 
 
 
 
 
 
 
 
 
905d227
 
3925884
6e12956
ca38a58
9255bb8
 
 
 
ad9b4e2
 
a47d2f3
ad9b4e2
9255bb8
 
 
 
 
 
 
 
 
6150c8f
83b80a1
6150c8f
 
 
 
 
83b80a1
adbf59b
 
6150c8f
 
 
 
6e12956
a47d2f3
 
c6b78e0
6e12956
a47d2f3
6e12956
 
 
 
 
 
a47d2f3
6e12956
6150c8f
 
c6b78e0
 
6150c8f
 
 
 
a47d2f3
6e12956
905d227
6e12956
 
 
 
 
 
 
 
a47d2f3
6e12956
a47d2f3
9255bb8
6150c8f
 
3e1163e
a47d2f3
 
6150c8f
 
cd7fced
a47d2f3
3925884
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a47d2f3
6150c8f
cd7fced
 
 
6150c8f
 
 
9255bb8
 
 
 
 
 
 
199759c
9255bb8
6150c8f
 
6e12956
b2ae2db
6150c8f
6e12956
7e9e544
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
import json
import os
import shutil
import requests
import spaces
import torch

import gradio as gr
from huggingface_hub import Repository
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css

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

HF_TOKEN = os.environ.get("HF_TOKEN", None)
CHECKPOINT_URL = "Salesforce/codegen-350M-mono"

SQLMODEL_PREFIX_URL = "luna-code/sqlmodel-codegen-350M-mono-prefix"
SFEPY_PREFIX_URL = "luna-code/sfepy-codegen-350M-mono-prefix"
MEGENGINE_PREFIX_URL = "luna-code/megengine-codegen-350M-mono-prefix"
MAIN_EVO_PREFIX_URL = "luna-code/codegen-350M-mono-evo-prefix"

SQLMODEL_FFT_URL = "luna-code/sqlmodel-codegen-350M-mono-fft"
SFEPY_FFT_URL = "luna-code/sfepy-codegen-350M-mono-fft"
MEGENGINE_FFT_URL = "luna-code/megengine-codegen-350M-mono-fft"
MAIN_EVO_FFT_URL = "luna-code/codegen-350M-mono-evo-fft"
MAIN_FD_FFT_URL = "luna-code/codegen-350M-mono-fd-fft"

LANGCHAIN_PREFIX_URL = "luna-code/langchain-codegen-350M-mono-prefix"
LLAMAINDEX_PREFIX_URL = "luna-code/llamaindex-codegen-350M-mono-prefix"
DSPY_PREFIX_URL = "luna-code/dspy-codegen-350M-mono-prefix"
CS_EVO_PREFIX_URL = "luna-code/cs-codegen-350M-mono-evo-prefix"

tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_URL)
basemodel = AutoModelForCausalLM.from_pretrained(CHECKPOINT_URL, device_map="auto")

sql_prefix = PeftModel.from_pretrained(basemodel, SQLMODEL_PREFIX_URL, device_map="auto")
sfepy_prefix = PeftModel.from_pretrained(basemodel, SFEPY_PREFIX_URL, device_map="auto")
megengine_prefix = PeftModel.from_pretrained(basemodel, MEGENGINE_PREFIX_URL, device_map="auto")
main_evo_prefix = PeftModel.from_pretrained(basemodel, MAIN_EVO_PREFIX_URL, device_map="auto")

sqlmodel_fft = AutoModelForCausalLM.from_pretrained(SQLMODEL_FFT_URL, device_map="auto")
sfepy_fft = AutoModelForCausalLM.from_pretrained(SFEPY_FFT_URL, device_map="auto")
megengine_fft = AutoModelForCausalLM.from_pretrained(MEGENGINE_FFT_URL, device_map="auto")
main_evo_fft = AutoModelForCausalLM.from_pretrained(MAIN_EVO_FFT_URL, device_map="auto")
main_fd_fft = AutoModelForCausalLM.from_pretrained(MAIN_FD_FFT_URL, device_map="auto")

langchain_prefix = PeftModel.from_pretrained(basemodel, LANGCHAIN_PREFIX_URL, device_map="auto")
llamaindex_prefix = PeftModel.from_pretrained(basemodel, LLAMAINDEX_PREFIX_URL, device_map="auto")
dspy_prefix = PeftModel.from_pretrained(basemodel, DSPY_PREFIX_URL, device_map="auto")
cs_evo_prefix = PeftModel.from_pretrained(basemodel, CS_EVO_PREFIX_URL, device_map="auto")

# basemodel = ""
# sql_prefix = ""
# sfepy_prefix = ""
# megengine_prefix = ""
# main_evo_prefix = ""
# sqlmodel_fft = ""
# sfepy_fft = ""
# megengine_fft = ""
# main_evo_fft = ""
# main_fd_fft = ""
# langchain_prefix = ""
# llamaindex_prefix = ""
# dspy_prefix = ""
# cs_evo_prefix = ""


model_map = {
    "Base": basemodel,
    "SQLModel Prefix": sql_prefix,
    "SfePy Prefix": sfepy_prefix,
    "MegEngine Prefix": megengine_prefix,
    "Main Evo Prefix": main_evo_prefix,
    "SQLModel FFT": sqlmodel_fft,
    "SfePy FFT": sfepy_fft,
    "MegEngine FFT": megengine_fft,
    "Main Evo FFT": main_evo_fft,
    "Main FD FFT": main_fd_fft,
    "LangChain Prefix": langchain_prefix,
    "LlamaIndex Prefix": llamaindex_prefix,
    "DSpy Prefix": dspy_prefix,
    "CS Evo Prefix": cs_evo_prefix,
}


FIM_PREFIX = "<fim_prefix>"
FIM_MIDDLE = "<fim_middle>"
FIM_SUFFIX = "<fim_suffix>"

FIM_INDICATOR = "<FILL_HERE>"

FORMATS = """## Model Formats

The model is pretrained on code and is formatted with special tokens in addition to the pure code data,\
such as prefixes specifying the source of the file or tokens separating code from a commit message.\
Use these templates to explore the model's capacities:

### 1. Prefixes 🏷️
For pure code files, use any combination of the following prefixes:

```
<reponame>REPONAME<filename>FILENAME<gh_stars>STARS\ncode<|endoftext|>
```
STARS can be one of: 0, 1-10, 10-100, 100-1000, 1000+

### 2. Commits 💾
The commits data is formatted as follows:

```
<commit_before>code<commit_msg>text<commit_after>code<|endoftext|>
```

### 3. Jupyter Notebooks 📓
The model is trained on Jupyter notebooks as Python scripts and structured formats like:

```
<start_jupyter><jupyter_text>text<jupyter_code>code<jupyter_output>output<jupyter_text>
```

### 4. Issues 🐛
We also trained on GitHub issues using the following formatting:
```
<issue_start><issue_comment>text<issue_comment>...<issue_closed>
```

### 5. Fill-in-the-middle 🧩
Fill in the middle requires rearranging the model inputs. The playground handles this for you - all you need is to specify where to fill:
```
code before<FILL_HERE>code after
```
"""

theme = gr.themes.Monochrome(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    radius_size=gr.themes.sizes.radius_sm,
    font=[
        gr.themes.GoogleFont("Open Sans"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    ],
)

def stream(model, code, generate_kwargs):
    input_ids = tokenizer(code, return_tensors="pt").to(device)
    generated_ids = model.generate(**input_ids, **generate_kwargs)
    return tokenizer.decode(generated_ids[0][input_ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()

@spaces.GPU(enable_queue=True)
def generate(
    prompt, temperature=0.6, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0, library="LangChain", method="Prefix"
):

    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )

    if method == "Base":
        output = stream(basemodel, prompt, generate_kwargs)
    elif method == "Prefix":
        output = stream(model_map[library + " Prefix"], prompt, generate_kwargs)
    elif method == "Evo Prefix" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map["Main Evo Prefix"], prompt, generate_kwargs)
    elif method == "FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map[library + " FFT"], prompt, generate_kwargs)
    elif method == "Evo FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map["Main Evo FFT"], prompt, generate_kwargs)
    elif method == "Full Data FFT" and library in ["SQLModel", "SfePy", "MegEngine"]:
        output = stream(model_map["Main FD FFT"], prompt, generate_kwargs)
    elif method == "Evo Prefix" and library in ["LangChain", "LlamaIndex", "DSPy"]:
        output = stream(model_map["CS Evo Prefix"], prompt, generate_kwargs)
    else:
        output = ""

    return output


examples = [
    "X_train, y_train, X_test, y_test = train_test_split(X, y, test_size=0.1)\n\n# Train a logistic regression model, predict the labels on the test set and compute the accuracy score",
    "// Returns every other value in the array as a new array.\nfunction everyOther(arr) {",
    "Poor English: She no went to the market. Corrected English:",
    "def alternating(list1, list2):\n   results = []\n   for i in range(min(len(list1), len(list2))):\n       results.append(list1[i])\n       results.append(list2[i])\n   if len(list1) > len(list2):\n       <FILL_HERE>\n   else:\n       results.extend(list2[i+1:])\n   return results",
]


def process_example(args):
    for x in generate(args):
        pass
    return x


css = ".generating {visibility: hidden}"

monospace_css = """
#q-input textarea {
    font-family: monospace, 'Consolas', Courier, monospace;
}
"""


css += share_btn_css + monospace_css + ".gradio-container {color: black}"


description = """
<div style="text-align: center;">
    <h1> 🌙 LUNA Models Playground</h1>
</div>
<div style="text-align: left;">
    <p>This is a demo to generate text and code with unknown libraries. The supported based model is <a href="Salesforce/codegen-350M-mono" style='color: #e6b800;'>CodeGen-350M-mono</a></p>
    <p>The supported libraries are:</p>
    <ul>
        <li><a href="https://sqlmodel.tiangolo.com" style='color: #e6b800;'>SQLModel</a></li>
        <li><a href="https://sfepy.org" style='color: #e6b800;'>SfePy</a></li>
        <li><a href="https://megengine.org" style='color: #e6b800;'>MegEngine</a></li>
        <li><a href="https://www.langchain.com/" style='color: #e6b800;'>LangChain</a></li>
        <li><a href="https://www.llamaindex.ai/" style='color: #e6b800;'>LlamaIndex</a></li>
        <li><a href="https://dspy-docs.vercel.app/" style='color: #e6b800;'>DSpy</a></li>
    </ul>
    <p><b>Please note:</b> These models are not designed for instruction purposes.</p>
</div>
"""
disclaimer = """⚠️<b>Any use or sharing of this demo constitues your acceptance of the BigCode [OpenRAIL-M](spaces/bigcode/bigcode-model-license-agreement) License Agreement and the use restrictions included within.</b>\
 <br>**Intended Use**: this app and its [supporting model](bigcode) are provided for demonstration purposes; not to serve as replacement for human expertise. For more details on the model's limitations in terms of factuality and biases, see the [model card.](hf.co/bigcode)"""

with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:
    with gr.Column():
        gr.Markdown(description)
        with gr.Row():
            library = gr.Dropdown(
                        ["SQLModel", "SfePy", "MegEngine", "LangChain", "LlamaIndex", "DSPy"],
                        value="LangChain",
                        label="Library",
                        info="Choose a library from the list",
                        )
        with gr.Row():
            method = gr.Dropdown(
                        ["Base", "Prefix", "Evo Prefix", "FFT", "Evo FFT", "Full Data FFT"],
                        value="Prefix",
                        label="Model",
                        info="Choose an expert from the list",
                        )
        with gr.Row():
            with gr.Column():
                instruction = gr.Textbox(
                    placeholder="Enter your code here",
                    lines=5,
                    label="Input",
                    elem_id="q-input",
                )
                submit = gr.Button("Generate", variant="primary")
                output = gr.Code(elem_id="q-output", lines=30, label="Output")
                with gr.Row():
                    with gr.Column():
                        with gr.Accordion("Advanced settings", open=False):
                            with gr.Row():
                                column_1, column_2 = gr.Column(), gr.Column()
                                with column_1:
                                    temperature = gr.Slider(
                                        label="Temperature",
                                        value=0.2,
                                        minimum=0.0,
                                        maximum=1.0,
                                        step=0.05,
                                        interactive=True,
                                        info="Higher values produce more diverse outputs",
                                    )
                                    max_new_tokens = gr.Slider(
                                        label="Max new tokens",
                                        value=256,
                                        minimum=0,
                                        maximum=8192,
                                        step=64,
                                        interactive=True,
                                        info="The maximum numbers of new tokens",
                                    )
                                with column_2:
                                    top_p = gr.Slider(
                                        label="Top-p (nucleus sampling)",
                                        value=0.90,
                                        minimum=0.0,
                                        maximum=1,
                                        step=0.05,
                                        interactive=True,
                                        info="Higher values sample more low-probability tokens",
                                    )
                                    repetition_penalty = gr.Slider(
                                        label="Repetition penalty",
                                        value=1.2,
                                        minimum=1.0,
                                        maximum=2.0,
                                        step=0.05,
                                        interactive=True,
                                        info="Penalize repeated tokens",
                                    )
                                    
                gr.Markdown(disclaimer)
                with gr.Group(elem_id="share-btn-container"):
                    community_icon = gr.HTML(community_icon_html, visible=True)
                    loading_icon = gr.HTML(loading_icon_html, visible=True)
                    share_button = gr.Button(
                        "Share to community", elem_id="share-btn", visible=True
                    )
                gr.Examples(
                    examples=examples,
                    inputs=[instruction],
                    cache_examples=False,
                    fn=process_example,
                    outputs=[output],
                )
                gr.Markdown(FORMATS)

    submit.click(
        generate,
        inputs=[instruction, temperature, max_new_tokens, top_p, repetition_penalty, library, method],
        outputs=[output]
    )
    share_button.click(None, [], [])
demo.queue().launch(debug=True)