File size: 7,780 Bytes
4375b7f
4e683ec
76a154f
 
 
b1c12fa
76a154f
c8fce50
4e683ec
cea1140
33012a5
cea1140
 
 
 
 
33012a5
13f3c6f
 
4375b7f
76a154f
4e683ec
c90359c
76a154f
 
25b70aa
76a154f
f5d5358
76a154f
 
 
4e683ec
98df5b4
76a154f
4e683ec
cea1140
 
b003c18
4e683ec
904cc64
f87f20f
625d8f8
904cc64
9938eb2
76a154f
 
4e683ec
76a154f
4e683ec
e1ed375
 
 
 
4e683ec
 
 
 
 
 
 
 
f87f20f
 
 
904cc64
f87f20f
 
 
 
 
 
 
594954c
 
 
 
 
4e683ec
9206e87
f87f20f
 
 
 
 
 
 
 
594954c
 
 
f424167
 
f87f20f
 
 
76a154f
f87f20f
 
 
 
76a154f
c18eb6d
0af44a3
ee9dd45
 
 
 
5084686
 
8736cd2
ee9dd45
 
0af44a3
ee9dd45
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0af44a3
ee9dd45
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5084686
 
6f2c8ad
0af44a3
ee9dd45
 
 
 
 
 
 
 
 
 
 
 
 
5084686
 
c18eb6d
e12823a
ee9dd45
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5084686
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e1ed375
5084686
 
 
 
 
 
e1ed375
5084686
 
 
 
f99f30f
5084686
 
 
 
4e683ec
 
 
7bcbb15
c3e546c
02c994c
25b70aa
76a154f
4e683ec
 
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
import os
from threading import Thread
from typing import Iterator

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

#import subprocess
# Install flash attention, skipping CUDA build if necessary
#subprocess.run(
#    "pip install flash-attn --no-build-isolation",
#    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
#    shell=True,
#)

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

DESCRIPTION = """\
# Try Patched Chat
"""

LICENSE = """\
---
This space was created by [patched](https://patched.codes).
"""

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 = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"
    model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True,trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.padding_side = 'right'
    # pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
    # tokenizer.use_default_system_prompt = FalseQwen/CodeQwen1.5-7B-Chat
    
@spaces.GPU(duration=60)
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int = 1024,
    temperature: float = 0.2,
    top_p: float = 0.95,
    # top_k: int = 50,
    # repetition_penalty: float = 1.2,
) -> Iterator[str]:
    conversation = []
    if system_prompt:
        conversation.append({"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})

#    prompt = pipe.tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
#    outputs = pipe(prompt, max_new_tokens=max_new_tokens, do_sample=True, temperature=temperature, top_p=top_p, 
#                   eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)

#    return outputs[0]['generated_text'][len(prompt):].strip()
    
    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)
    
    terminators = [
    tokenizer.eos_token_id,
    tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        #top_k=top_k,
        temperature=temperature,
        eos_token_id=terminators,
        #eos_token_id=tokenizer.eos_token_id, 
        #pad_token_id=tokenizer.pad_token_id,
        #num_beams=1,
        #repetition_penalty=1.2,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)

example1='''Fix vulnerability CWE-327: Use of a Broken or Risky Cryptographic Algorithm in the following code snippet.

def md5_hash(path):
     with open(path, "rb") as f:
         content = f.read()
     return hashlib.md5(content).hexdigest()
'''

example2='''Carefully analyze the given old code and new code and generate a summary of the changes.

Old Code:

#include <stdio.h>
#include <stdlib.h>

typedef struct Node {
    int data;
    struct Node *next;
} Node;

void processList() {
    Node *head = (Node*)malloc(sizeof(Node));
    head->data = 1;
    head->next = (Node*)malloc(sizeof(Node));
    head->next->data = 2;

    printf("First element: %d\n", head->data);

    free(head->next); 
    free(head); 

    printf("Accessing freed list: %d\n", head->next->data);
}

New Code:

#include <stdio.h>
#include <stdlib.h>

typedef struct Node {
    int data;
    struct Node *next;
} Node;

void processList() {
    Node *head = (Node*)malloc(sizeof(Node));
    if (head == NULL) {
        perror("Failed to allocate memory for head");
        return;
    }

    head->data = 1;
    head->next = (Node*)malloc(sizeof(Node));
    if (head->next == NULL) {
        free(head);
        perror("Failed to allocate memory for next node");
        return;
    }
    head->next->data = 2;

    printf("First element: %d\n", head->data);

    free(head->next); 
    head->next = NULL; 
    free(head); 
    head = NULL; 

    if (head != NULL && head->next != NULL) {
        printf("Accessing freed list: %d\n", head->next->data);
    }
}
'''

example3='''Is the following code prone to CWE-117: Improper Output Neutralization for Logs. Respond only with YES or NO.

from flask import Flask, request, jsonify
import logging

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@app.route('/api/data', methods=['GET'])
def get_data():
    api_key = request.args.get('api_key')
    logger.info("Received request with API Key: %s", api_key)  
    data = {"message": "Data processed"}
    return jsonify(data)
'''

example4='''Fix vulnerability CWE-78: Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection') in the following code snippet.

def run(command, desc=None, errdesc=None, custom_env=None, live: bool = default_command_live) -> str:
    if desc is not None:
        print(desc)
    run_kwargs = {{
        "args": command,
        "shell": True,
        "env": os.environ if custom_env is None else custom_env,
        "encoding": 'utf8',
        "errors": 'ignore',
    }}
    if not live:
        run_kwargs["stdout"] = run_kwargs["stderr"] = subprocess.PIPE
    result = subprocess.run(**run_kwargs)  ##here
    if result.returncode != 0:
        error_bits = [
            f"{{errdesc or 'Error running command'}}.",
            f"Command: {{command}}",
            f"Error code: {{result.returncode}}",
        ]
        if result.stdout:
            error_bits.append(f"stdout: {{result.stdout}}")
        if result.stderr:
            error_bits.append(f"stderr: {{result.stderr}}")
        raise RuntimeError("\n".join(error_bits))
    return (result.stdout or "")
'''

chat_interface = gr.ChatInterface(
    fn=generate,
    chatbot=gr.Chatbot(height="480px"),
    additional_inputs=[
        gr.Textbox(label="System prompt", lines=4),
        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.2,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.95,
        ),
    ],
    stop_btn=None,
    examples=[
        ["You are a helpful coding assistant. Create a snake game in Python."],
        [example1],
        [example2],
        [example3],
        [example4],
    ],
)

with gr.Blocks(css="style.css",) as demo:
    gr.Markdown(DESCRIPTION)
    chat_interface.render()
    gr.Markdown(LICENSE, elem_classes="contain")

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