Create app.py

app.py

@@ -0,0 +1,255 @@
+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
+
+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 = "meta-llama/Meta-Llama-3-8B-Instruct"
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=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()