Duplicate from facebook/incoder-demo

Co-authored-by: Daniel Fried <[email protected]>

.gitattributes

@@ -0,0 +1,28 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+debug
+unlock
+__pycache__
+incoder-6B

README.md

@@ -0,0 +1,15 @@
+---
+title: Incoder Demo
+emoji: 💻
+colorFrom: red
+colorTo: purple
+sdk: gradio
+sdk_version: 2.9.1
+python_version: 3.8.13
+app_file: start.py
+license: cc-by-nc-4.0
+pinned: false
+duplicated_from: facebook/incoder-demo
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

index.html

@@ -0,0 +1 @@
+demo is loading

modules/app.py

@@ -0,0 +1,240 @@
+import sys
+from typing import List
+import traceback
+import os
+import base64
+
+import logging
+logging.basicConfig(level=logging.INFO)
+import modules.cloud_logging
+
+import tokenizers
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import json
+import pprint
+
+# needs to be imported *before* transformers
+if os.path.exists('debug'):
+    BIG_MODEL = False
+    CUDA = False
+else:
+    BIG_MODEL = True
+    CUDA = True
+
+# from flask import Flask, request, render_template
+# from flask_cors import CORS
+# app = Flask(__name__, static_folder='static')
+# app.config['TEMPLATES_AUTO_RELOAD'] = Tru
+# CORS(app, resources= {
+#     r"/generate": {"origins": origins},
+#     r"/infill": {"origins": origins},
+# })
+# origins=[f"http://localhost:{PORT}", "https://huggingface.co", "https://hf.space"]
+
+PORT = 7860
+VERBOSE = False
+
+if os.path.exists('unlock'):
+    MAX_LENGTH = 2048
+else:
+    MAX_LENGTH = 256+64
+TRUNCATION_MESSAGE = f'warning: This demo is limited to {MAX_LENGTH} tokens in the document for efficiency.'
+
+if BIG_MODEL:
+    model_name = "facebook/incoder-6B"
+    kwargs = dict(
+        revision="float16",
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+    )
+else:
+    model_name = "facebook/incoder-1B"
+    kwargs = dict()
+
+from fastapi import FastAPI, Request
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse, StreamingResponse
+app = FastAPI(docs_url=None, redoc_url=None)
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+
+logging.info("loading model")
+model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
+logging.info("loading tokenizer")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+logging.info("loading complete")
+
+if CUDA:
+    model = model.half().cuda()
+
+BOS = "<|endoftext|>"
+EOM = "<|endofmask|>"
+
+def make_sentinel(i):
+    return f"<|mask:{i}|>"
+
+SPECIAL_TOKENS = [make_sentinel(i) for i in range(256)] + [EOM]
+
+def generate(input, length_limit=None, temperature=None):
+    input_ids = tokenizer(input, return_tensors="pt").input_ids
+    if CUDA:
+        input_ids = input_ids.cuda()
+    current_length = input_ids.flatten().size(0)
+    max_length = length_limit + current_length
+    truncated = False
+    if max_length > MAX_LENGTH:
+        max_length = MAX_LENGTH
+        truncated = True
+    if max_length == current_length:
+        return input, True
+    output = model.generate(input_ids=input_ids, do_sample=True, top_p=0.95, temperature=temperature, max_length=max_length)
+    detok_hypo_str = tokenizer.decode(output.flatten())
+    if detok_hypo_str.startswith(BOS):
+        detok_hypo_str = detok_hypo_str[len(BOS):]
+    return detok_hypo_str, truncated
+
+def infill(parts: List[str], length_limit=None, temperature=None, extra_sentinel=False, max_retries=1):
+    assert isinstance(parts, list)
+    retries_attempted = 0
+    done = False
+
+
+    while (not done) and (retries_attempted < max_retries):
+        any_truncated = False
+        retries_attempted += 1
+        if VERBOSE:
+            logging.info(f"retry {retries_attempted}")
+        if len(parts) == 1:
+            prompt = parts[0]
+        else:
+            prompt = ""
+            # encode parts separated by sentinel
+            for sentinel_ix, part in enumerate(parts):
+                prompt += part
+                if extra_sentinel or (sentinel_ix < len(parts) - 1):
+                    prompt += make_sentinel(sentinel_ix)
+            
+            # prompt += TokenizerWrapper.make_sentinel(0)
+        
+        infills = []
+        complete = []
+
+        done = True
+
+        for sentinel_ix, part in enumerate(parts[:-1]):
+            complete.append(part)
+            prompt += make_sentinel(sentinel_ix)
+            completion, this_truncated = generate(prompt, length_limit, temperature)
+            any_truncated |= this_truncated
+            completion = completion[len(prompt):]
+            if EOM not in completion:
+                if VERBOSE:
+                    logging.info(f"warning: {EOM} not found")
+                completion += EOM
+                # TODO: break inner loop here
+                done = False
+            completion = completion[:completion.index(EOM) + len(EOM)]
+            infilled = completion[:-len(EOM)]
+            infills.append(infilled)
+            complete.append(infilled)
+            prompt += completion
+        complete.append(parts[-1])
+        text = ''.join(complete)
+
+    if VERBOSE:
+        logging.info("generated text:")
+        logging.info(prompt)
+        logging.info()
+        logging.info("parts:")
+        logging.info(parts)
+        logging.info()
+        logging.info("infills:")
+        logging.info(infills)
+        logging.info()
+        logging.info("restitched text:")
+        logging.info(text)
+        logging.info()
+    
+    return {
+        'text': text,
+        'parts': parts,
+        'infills': infills,
+        'retries_attempted': retries_attempted,
+        'truncated': any_truncated,
+    } 
+
+
+@app.head("/")
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(path="static/index.html", media_type="text/html")
+
+@app.get('/generate')
+# async def generate_maybe(request: Request):
+async def generate_maybe(info: str):
+    # form = await info.json()
+    # form = await request.json()
+    # info is a base64-encoded, url-escaped json string (since GET doesn't support a body, and POST leads to CORS issues)
+    # fix padding, following https://stackoverflow.com/a/9956217/1319683
+    info = base64.urlsafe_b64decode(info + '=' * (4 - len(info) % 4)).decode('utf-8')
+    form = json.loads(info)
+    # print(form)
+    prompt = form['prompt']
+    length_limit = int(form['length'])
+    temperature = float(form['temperature'])
+    logging.info(json.dumps({
+        'length': length_limit,
+        'temperature': temperature,
+        'prompt': prompt,
+    }))
+    try:
+        generation, truncated = generate(prompt, length_limit, temperature)
+        if truncated:
+            message = TRUNCATION_MESSAGE 
+        else:
+            message = ''
+        return {'result': 'success', 'type': 'generate', 'prompt': prompt, 'text': generation, 'message': message}
+    except Exception as e:
+        traceback.print_exception(*sys.exc_info())
+        logging.error(e)
+        return {'result': 'error', 'type': 'generate', 'prompt': prompt, 'message': f'Error: {e}.'}
+
+@app.get('/infill')
+# async def infill_maybe(request: Request):
+async def infill_maybe(info: str):
+    # form = await info.json()
+    # form = await request.json()
+    # info is a base64-encoded, url-escaped json string (since GET doesn't support a body, and POST leads to CORS issues)
+    # fix padding, following https://stackoverflow.com/a/9956217/1319683
+    info = base64.urlsafe_b64decode(info + '=' * (4 - len(info) % 4)).decode('utf-8')
+    form = json.loads(info)
+    length_limit = int(form['length'])
+    temperature = float(form['temperature'])
+    max_retries = 1
+    extra_sentinel = True
+    logging.info(json.dumps({
+        'length': length_limit,
+        'temperature': temperature,
+        'parts_joined': '<infill>'.join(form['parts']),
+    }))
+    try:
+        if len(form['parts']) > 4:
+            return {'result': 'error', 'text': ''.join(form['parts']), 'type': 'infill', 'message': f"error: Can't use more than 3 <infill> tokens in this demo (for efficiency)."}
+        generation = infill(form['parts'], length_limit, temperature, extra_sentinel=extra_sentinel, max_retries=max_retries)
+        generation['result'] = 'success'
+        generation['type'] = 'infill'
+        if generation['truncated']:
+            generation['message'] = TRUNCATION_MESSAGE
+        else:
+            generation['message'] = ''
+        return generation
+        # return {'result': 'success', 'prefix': prefix, 'suffix': suffix,  'text': generation['text']}
+    except Exception as e:
+        traceback.print_exception(*sys.exc_info())
+        logging.error(e)
+        return {'result': 'error', 'type': 'infill', 'message': f'Error: {e}.'}
+
+
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=PORT, threaded=False)

modules/cloud_logging.py

@@ -0,0 +1,21 @@
+import os
+def make_logging_client():
+    cred_filename = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')
+    if not cred_filename:
+        return None
+    print("cred filename:", cred_filename)
+    cred_string = os.environ.get('GOOGLE_APPLICATION_CREDENTIALS_STRING')
+    print("cred string:", bool(cred_string))
+    if not os.path.exists(cred_filename):
+        if cred_string:
+            print(f"writing cred string to {cred_filename}")
+            with open(cred_filename, 'w') as f:
+                f.write(cred_string)
+        else:
+            return None
+    from google.cloud import logging
+    logging_client = logging.Client()
+    logging_client.setup_logging()
+    return logging_client
+
+logging_client = make_logging_client()

packages.txt

@@ -0,0 +1 @@
+rustc

requirements.txt

@@ -0,0 +1,7 @@
+fastapi==0.74.*
+requests==2.27.*
+torch==1.11.*
+uvicorn[standard]==0.17.*
+tokenizers==0.12.1
+git+https://github.com/huggingface/transformers.git@b18dfd95e1f60ae65a959a7b255fc06522170d1b
+google-cloud-logging

start.py

@@ -0,0 +1,3 @@
+import subprocess
+
+subprocess.run("uvicorn modules.app:app --timeout-keep-alive 300 --host 0.0.0.0 --port 7860", shell=True)

static/frame.html

@@ -0,0 +1 @@
+<iframe src="index.html"></iframe>

static/index.html

@@ -0,0 +1,633 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <meta charset="utf-8"/>
+    <meta name="viewport" contents="width=device-width, initial-scale=1.0" />
+    <title>InCoder</title>
+    <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.4.1/jquery.min.js"></script>
+    <script src="https://cdn.jsdelivr.net/npm/[email protected]/base64.min.js"></script>
+
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/ace.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-plain_text.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-c_cpp.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-csharp.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-clojure.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-coffee.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-golang.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-haskell.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-python.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-java.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-javascript.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-lua.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-objectivec.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-perl.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-php.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-python.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-ruby.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-rust.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-scala.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-sh.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-swift.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/ace/1.4.14/mode-typescript.min.js"></script>
+    <link rel="stylesheet" href="static/style.css">
+</head>
+<style type="text/css">
+/* body {
+    font-family: sans-serif;
+} */
+/* .leftside {
+} */
+main {
+  max-width: 80rem;
+}
+.rightside {
+    width: 30em;
+}
+.submit-holder {
+    margin-top: 2em;
+}
+.submit input {
+    font-size: 16pt;
+}
+.slider {
+    width: 20em;
+}
+#faq {
+    max-width: 60em;
+}
+#result {
+    font-family: monospace;
+    white-space: pre-wrap;
+    word-wrap: break-word;
+    font-size: 12pt;
+    clear: both;
+    margin-top: 1em;
+    border: 1px solid black;
+    padding: 1em;
+    width: 60em;
+    min-height: 12em;
+}
+#prompt {
+    font-weight: bold;
+}
+.loader {
+    border: 4px solid #f3f3f3;
+    border-radius: 50%;
+    border-top: 4px solid #3498db;
+    width: 30px;
+    height: 30px;
+    animation: spin 2s linear infinite;
+    margin-right: 1em;
+}
+@keyframes spin {
+    0% { transform: rotate(0deg); }
+    100% { transform: rotate(360deg); }
+}
+#loader_holder {
+    visibility: hidden;
+    display: flex;
+    align-items: center;
+}
+
+label {
+    margin-top: 1em;
+    display: inline-elock;
+    width: 10em;
+    text-align: right;
+    font-size: 80%;
+}
+#loader_holder_super {
+}
+#error {
+    color: red;
+    width: 100%;
+}
+#warning {
+    color: darkorange;
+    width: 100%;
+}
+#examples span {
+    margin-right: 1em;
+}
+#editor {
+    position: relative;
+    width: 100%;
+    height: 400px;
+}
+#editor-holder {
+    position: relative;
+    width: 100%;
+    height: 400px;
+}
+.ace_infill {
+    color: red;
+}
+</style>
+<body>
+  <main>
+    <div class="card" id="about">
+        <div class="header"> <h1>InCoder: A Generative Model for Code Infilling and Synthesis</h1> </div>
+        <p>Demo of the 6.7B parameter version of InCoder: a decoder-only Transformer model that can both extend and insert/infill code.</p>
+        <p>Select one of the examples below, or input your own code into the editor. You can type &lt;infill&gt; to mark a location you want the model to insert code at.</p>
+        <p>Click "Extend" to append text at the end of the editor. Click "Infill" to replace all &lt;infill&gt; masks. (Click "Add &lt;infill&gt; mask" to add a mask at the cursor or replace the current selection.) </p>
+    </div>
+    <div class="card" id="examples">
+        <div id="examples-infill">
+            <span class="softspan">Infill Examples:</span>
+            <br>
+            <span class="softspan"><a href='javascript:select_example("type-pred");'>Type prediction</a></span>
+            <span class="softspan"><a href='javascript:select_example("multi-region");'>Docstring to function</a></span>
+            <span class="softspan"><a href='javascript:select_example("docstring-2");'>Function to docstring</a></span>
+            <!--
+            <span class="softspan"><a href='javascript:select_example("python-infill2");'>Docstring to function</a></span>
+            -->
+            <span class="softspan"><a href='javascript:select_example("class");'>Class generation</a></span>
+        </div>
+        <div id="examples-extend">
+            <span class="softspan">Extend Examples:</span>
+            <br>
+            <span class="softspan"><a href='javascript:select_example("python");'>Python</a></span>
+            <span class="softspan"><a href='javascript:select_example("javascript");'>JavaScript</a></span>
+            <span class="softspan"><a href='javascript:select_example("jupyter");'>Jupyter</a></span>
+            <span class="softspan"><a href='javascript:select_example("stackoverflow");'>StackOverflow</a></span>
+            <span class="softspan"><a href='javascript:select_example("metadata-conditioning");'>Metadata Conditioning</a></span>
+            <span class="softspan"><a href='javascript:select_example("metadata-prediction");'>Metadata Prediction</a></span>
+        </div>
+    </div>
+    <div class="card" id="controls">
+        <div>
+            <label>Num Tokens:</label>
+            <input type="range" value="64" min="16" max="256" step="16" class="slider"
+                oninput="this.nextElementSibling.value = this.value" name="length" id='length_slider'>
+            <output class='a' id="length_slider_output">64</output>
+        </div>
+        <div>
+            <label>Temperature:</label>
+            <input type="range" value="0.6" min="0.1" max="1.0" step="0.10" class="slider"
+                oninput="this.nextElementSibling.value = this.value" name="temp" id='temp_slider'>
+            <output class='a' id="temp_slider_output">0.6</output>
+        </div>
+        <div id="buttons">
+          <br>
+          <input type="button" value="Extend" id="extend-form-button" />
+          <input type="button" value="Infill" id="infill-form-button" />
+          <br>
+          <br>
+          <input type="button" value="Add <infill> mask" id="insert-mask-button" title="add the infill marker at cursor or selection" />
+        </div>
+    </div>
+    <div id="edit-container" class="card">
+        <div id="syntax">
+            <span class="softspan">Syntax:</span>
+            <select name="mode" id="mode">
+                <option value="text">Text</option>
+                <option value="c_cpp">C/C++</option>
+                <option value="csharp">C#</option>
+                <option value="clojure">Clojure</option>
+                <option value="coffee">CoffeeScript</option>
+                <option value="golang">Go</option>
+                <option value="haskell">Haskell</option>
+                <option value="java">Java</option>
+                <option value="javascript">JavaScript</option>
+                <option value="lua">Lua</option>
+                <option value="objectivec">Objective C</option>
+                <option value="perl">Perl</option>
+                <option value="php">PHP</option>
+                <option value="python">Python</option>
+                <option value="ruby">Ruby</option>
+                <option value="rust">Rust</option>
+                <option value="scala">Scala</option>
+                <option value="sh">Shell</option>
+                <option value="swift">Swift</option>
+                <option value="typescript">Typescript</option>
+            </select>
+        </div>
+        <div id="editor"></div>
+    </div>
+    <div id="loader_holder_super" class="card">
+        <h1>Messages</h1>
+        <div id="error"></div>
+        <div id="warning"></div>
+        <div id="loader_holder">
+            <div class="loader"></div>
+            <div>
+                Generation queued, please wait...
+            </div>
+        </div>
+    </div>
+    <div id="info" class="card">
+        <h1 id="debug-info">More Info</h3>
+        <p>
+            See <a href="https://sites.google.com/view/incoder-code-models" target="_blank" rel="noopener noreferrer">our project site</a> for more information on
+            these models, including a paper and examples.
+        </p>
+    
+        <p>
+            For instructions on setting up and using the models (via HuggingFace transformers), see 
+            <a href="https://github.com/dpfried/incoder/blob/main/README.md" target="_blank" rel="noopener noreferrer">our readme</a>.
+        </p>
+    
+        <h1 id="debug-info">Credits</h3>
+        <p>This model was developed at Facebook AI Research by Daniel Fried*, Armen Aghajanyan*, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong,
+            Wen-tau Yih, Luke Zettlemoyer, and Mike Lewis.</p>
+        <p>Thanks to Naman Goyal and Stephen Roller for writing the code this demo was based on. Extensions by Daniel Fried and
+            Sida Wang.</p>
+    </div>
+  </main>
+<script type="text/javascript">
+// these constants are only used for providing user expectations.
+var OVERHEAD = 3;
+var PER_TOKEN = 0.12;
+var SPLIT_TOKEN = "<infill>"
+
+var Range = require("ace/range").Range;
+
+// examples for the user
+var EXAMPLES = {
+    "python-infill2": {
+        "prompt":
+`<| file ext=.py |>
+from collections import Counter
+
+def <infill>
+    """Count the number of occurrences of each word in the file."""
+    <infill>
+`,
+        "length": 64,
+        "temperature": 0.2,
+        "mode": "python"
+    },
+    "multi-region": {
+        "prompt":
+`<| file ext=.py |>
+<infill>
+    """ Load the given gzip jsonl file. """
+    <infill>
+`,
+        "length": 64,
+        "temperature": 0.2,
+        "mode": "python"
+    },
+    "type-pred": {
+        "prompt": 
+`def count_words(filename: str) -> <infill>
+    """Count the number of occurrences of each word in the file."""
+    with open(filename, 'r') as f:
+        word_counts = {}
+        for line in f:
+            for word in line.split():
+                if word in word_counts:
+                    word_counts[word]  = 1
+                else:
+                    word_counts[word] = 1
+    return word_counts
+`,
+        "length": 4,
+        "temperature": 0.2,
+        "mode": "python"
+    },
+    "docstring-2": {
+        "prompt":
+`def _minimize_in_graph(build_loss_fn, num_steps=200, optimizer=None):
+  """
+  <infill>
+  """
+  optimizer = tf.compat.v1.train.AdamOptimizer(
+      0.1) if optimizer is None else optimizer
+
+  def train_loop_body(step):
+    train_op = optimizer.minimize(
+        build_loss_fn if tf.executing_eagerly() else build_loss_fn())
+    return tf.tuple(tensors=[tf.add(step, 1)], control_inputs=[train_op])
+
+  minimize_op = tf.compat.v1.while_loop(
+      cond=lambda step: step < num_steps,
+      body=train_loop_body,
+      loop_vars=[tf.constant(0)],
+      return_same_structure=True)[0] 
+  return minimize_op`,
+  "length": 64,
+  "temperature": 0.3,
+  "mode": "python",
+    },
+    "docstring": {
+        "prompt": 
+`<| file ext=.py |>
+
+def count_words(filename: str) -> Dict[str, int]:
+    """<infill>
+    """
+    with open(filename, 'r') as f:
+        word_counts = {}
+        for line in f:
+            for word in line.split():
+                if word in word_counts:
+                    word_counts[word]  = 1
+                else:
+                    word_counts[word] = 1
+    return word_counts
+`,
+        "length": 32,
+        "temperature": 0.2,
+        "mode": "python"
+    },
+    "python": {
+        "prompt": 
+`<| file ext=.py |>
+def count_words(filename):
+    """Count the number of occurrences of each word in the file"""`,
+        "length": 64,
+        "temperature": 0.6,
+        "mode": "python"
+    },
+    "class": {
+        "prompt": "<| file ext=.py |>\nclass Person:\n" + SPLIT_TOKEN + "\np = Person('Eren', 18, 'Male')",
+        "length": 64,
+        "temperature": 0.2,
+        "mode": "python"
+    },
+    "javascript": {
+        "prompt": "// fetch from the given URL and load the response contents into a new div",
+        "length": 64,
+        "temperature": 0.6,
+        "mode": "javascript"
+    },
+    "jupyter": {
+        "prompt": "<| file ext=.ipynb:python |>\n<text>\nThis notebook demonstrates using scikit-learn to perform PCA.\n</text>\n<cell>",
+        "length": 64,
+        "temperature": 0.6,
+        "mode": "python"
+    },
+    "stackoverflow": {
+        "prompt": "<| q tags=regex,html |>\nParsing HTML with regular expressions\nHow do I do this? Is it a good idea?\n<|/ q dscore=3 |>\n<| a dscore=4 |>",
+        "length": 64,
+        "temperature": 0.6,
+        "mode": "text"
+    },
+    "metadata-conditioning": {
+        "prompt": "<| file ext=.py filename=train_model.py source=github dstars=4 |>\n",
+        "length": 64,
+        "temperature": 0.6,
+        "mode": "python"
+    },
+    "metadata-prediction": {
+        "prompt": "<| file source=github ext=.py |>\nfrom setuptools import setup\nfrom setuptools_rust import Binding, RustExtension\n\nextras = {}\nextras[\"testing\"] = [\"pytest\", \"requests\", \"numpy\", \"datasets\"]\nextras[\"docs\"] = [\"sphinx\", \"sphinx_rtd_theme\", \"setuptools_rust\"]\n\nsetup(\n    name=\"tokenizers\",\n    version=\"0.11\",\n    description=\"Fast and Customizable Tokenizers\",\n    long_description=open(\"README.md\", \"r\", encoding=\"utf-8\").read(),\n)\n\n<|/ file filename=",
+        "length": 1,
+        "temperature": 0.2,
+        "mode": "python"
+    },
+    "humaneval": {
+        "prompt": "from typing import List, Optional\n\n\ndef longest(strings: List[str]) -> Optional[str]:\n    \"\"\" Out of list of strings, return the longest one. Return the first one in case of multiple\n    strings of the same length. Return None in case the input list is empty.\n    >>> longest([])\n\n    >>> longest(['a', 'b', 'c'])\n    'a'\n    >>> longest(['a', 'bb', 'ccc'])\n    'ccc'\n    \"\"\"\n",
+        "temperature": 0.6,
+        "length": 64,
+        "mode": "python"
+    },
+};
+
+var editor = ace.edit("editor");
+editor.setOption("wrap", true);
+//var editor = null;
+
+function set_editor_mode(mode) {
+    session = editor.session
+    session.setMode("ace/mode/" + mode, function() {
+        var rules = session.$mode.$highlightRules.getRules();
+        for (var stateName in rules) {
+            if (Object.prototype.hasOwnProperty.call(rules, stateName)) {
+                rules[stateName].unshift({
+                    token: 'infill',
+                    regex: SPLIT_TOKEN 
+                });
+            }
+        }
+        // force recreation of tokenizer
+        session.$mode.$tokenizer = null;
+        session.bgTokenizer.setTokenizer(session.$mode.getTokenizer());
+        // force re-highlight whole document
+        session.bgTokenizer.start(0);
+    });
+}
+
+/*
+var textarea = $('textarea[name="prompt"]').hide();
+var prefix_textarea = $('textarea[name="prefix"]').hide();
+var suffix_textarea = $('textarea[name="suffix"]').hide();
+editor.getSession().on('change', function () {
+    textarea.val(editor.getSession().getValue());
+});
+*/
+
+function set_text(text) {
+    editor.getSession().setValue(text);
+    // textarea.val(text);
+}
+
+function set_selection(data) {
+    var lines = editor.getSession().doc.$lines;
+    var lines_flat = join_lines(lines);
+    if (data['type'] == 'generate') {
+        doc_length = lines_flat.length;
+        var start = convert_string_index_to_location(data['prompt'].length, lines);
+        var end = convert_string_index_to_location(doc_length, lines);
+        // reverse this so that we can shift select to shorten and delete extra stuff
+        editor.selection.setRange(new Range(end.row, end.column, start.row, start.column));
+    } else if (data['type'] == 'infill') {
+        var length_so_far = 0;
+        for (var i = 0; i < data['infills'].length; i++) {
+            var prefix = data['parts'][i];
+            var suffix = data['parts'][i+1];
+            var infilled = data['infills'][i];
+            var start = convert_string_index_to_location(length_so_far + prefix.length, lines);
+            var end = convert_string_index_to_location(length_so_far + (prefix + infilled).length, lines);
+            var range = null;
+            if (data['infills'].length == 1) {
+                range = new Range(end.row, end.column, start.row, start.column)
+            } else {
+                range = new Range(start.row, start.column, end.row, end.column)
+            }
+            if (i == 0) {
+                editor.selection.setRange(range);
+            } else {
+                editor.selection.addRange(range);
+            }
+            length_so_far += (prefix + infilled).length;
+        }
+    }
+    editor.focus();
+}
+
+function select_example(name) {
+    $("#length_slider").val(EXAMPLES[name]["length"]);
+    $("#length_slider_output").text(EXAMPLES[name]["length"]);
+    $("#temp_slider").val(EXAMPLES[name]["temperature"]);
+    $("#temp_slider_output").text(EXAMPLES[name]["temperature"]);
+    set_text(EXAMPLES[name]["prompt"])
+    var mode = EXAMPLES[name]["mode"];
+
+    set_editor_mode(mode);
+    $("#mode").val(mode).change();
+}
+
+function newline_character() {
+    return editor.getSession().doc.getNewLineCharacter();
+}
+
+function join_lines(lines) {
+    return lines.join(newline_character());
+}
+
+function get_prefix(location, lines) {
+    if (!(location.hasOwnProperty('row') && location.hasOwnProperty('column'))) { 
+        console.error("invalid location " + location);
+    }
+    if (location.row == 0) {
+        return lines[location.row].substring(0, location.column);
+    } else {
+        return join_lines(lines.slice(0, location.row)) + newline_character() + lines[location.row].substring(0, location.column);
+    }
+}
+
+function convert_location_to_string_index(location, lines) {
+    return get_prefix(location, lines).length;
+}
+
+function convert_string_index_to_location(string_index, lines) {
+    var column = 0;
+    var row = 0;
+    var char_count = 0;
+    var line_sep_length = editor.getSession().doc.getNewLineCharacter().length;
+    for (var i = 0; i < lines.length; i++) {
+        var line = lines[i];
+        var new_char_count = char_count + line.length + line_sep_length;
+        if (string_index < new_char_count) {
+            return {
+                'row': i,
+                'column': string_index - char_count,
+            }
+        }
+        char_count = new_char_count;
+    }
+    console.error("did not find index " + string_index + " in lines " + lines);
+    return null;
+}
+
+function get_infill_parts(warn_on_single) {
+    var lines = editor.getSession().doc.$lines;
+    var lines_flat = join_lines(lines);
+    parts = lines_flat.split(SPLIT_TOKEN)
+    if (warn_on_single && parts.length == 1) {
+        window.alert('There are no infill masks, add some <infill> masks before requesting an infill')
+    }
+    return parts
+}
+
+function insert_mask() {
+    if (editor.selection.ranges.length > 1) {
+        for (var i = 0; i < editor.selection.ranges.length; i++) {
+            console.log('range is', editor.selection.ranges[i])
+            editor.session.replace(editor.selection.ranges[i], SPLIT_TOKEN)
+        }
+    } else {
+        editor.session.replace(editor.selection.getRange(), SPLIT_TOKEN)
+    }
+}
+
+
+function make_generate_listener(url) {
+    return async function(event) {
+        var length = $("#length_slider").val();
+        var eta = PER_TOKEN * length  + OVERHEAD;
+        // $("#eta").text(eta);
+        // $("#infill-form-button").click(function (event) { console.log(editor.selection.getCursor()); });
+
+        // get temperature and response length parameters
+        var send_data = {
+            length: $("#length_slider").val(),
+            temperature: $("#temp_slider").val(),
+            extra_sentinel: $('#extra_sentinel_checkbox').is(":checked"),
+            max_retries: $('#max_retries_slider').val(),
+            parts: get_infill_parts(url == "infill"),
+            prompt: editor.getSession().getValue(),
+        }
+        console.log("send_data:");
+        console.log(send_data);
+
+        $("#loader_holder").css("visibility", "visible");
+        $("#extend-form-button").prop("disabled", true);
+        $("#infill-form-button").prop("disabled", true);
+        $("#error").text("");
+
+        function complete() {
+            $("#loader_holder").css("visibility", "hidden");
+            $("#extend-form-button").prop("disabled", false);
+            $("#infill-form-button").prop("disabled", false);
+        }
+
+        function success(receive_data) {
+            console.log("Response:");
+            console.log(receive_data);
+            if (receive_data["result"] == "success") {
+                console.log("success");
+                // $("#prompt").text(data["prompt"]);
+                // $("#response").text(data["text"]);
+                set_text(receive_data["text"]);
+                set_selection(receive_data);
+                $("#error").text("");
+                if (receive_data["message"] != "") {
+                    $("#warning").text(receive_data["message"]);
+                } else {
+                    $("#warning").text("");
+                }
+            } else {
+                console.log("error");
+                set_text(receive_data["text"])
+                $("#error").text(receive_data["message"]);
+            }
+        }
+
+        function error(err) {
+            console.log(err);
+            $("#error").text(err);
+        }
+
+        try {
+            var stringified = JSON.stringify(send_data);
+            // var encoded_data = encodeURIComponent(btoa(stringified));
+            var encoded_data = Base64.encodeURI(stringified);
+
+            const response = await fetch(`${url}?info=${encoded_data}`);
+            // const response = await fetch(`${url}` {
+            //     method: 'GET',
+            //     body: encoded_data,
+            // });
+            if (response.status >= 400) {
+                error(response.statusText);
+                console.log("here");
+                console.log(response.status);
+            } else {
+                response.json().then(success).catch(error).finally(complete);
+            }
+        } catch (e) {
+            error(e);
+        } finally {
+            complete();
+        }
+    }
+}
+
+// actual logic
+$(document).ready(function() {
+    $("#insert-mask-button").click(insert_mask);
+    $("#extend-form-button").click(make_generate_listener("generate"));
+    $("#infill-form-button").click(make_generate_listener("infill"));
+    $("#mode").change(function (e) {
+        var mode = $("#mode").val();
+        set_editor_mode(mode);
+    });
+    select_example("python")
+    // set_editor_mode("python");
+});
+</script>
+  <script src="https://cdnjs.cloudflare.com/ajax/libs/iframe-resizer/4.3.2/iframeResizer.contentWindow.min.js"></script>
+</body>
+</html>

static/style.css

@@ -0,0 +1,39 @@
+body {
+	padding: 2rem;
+	font-family: -apple-system, BlinkMacSystemFont, "Arial", sans-serif;
+}
+
+h1 {
+	font-size: 16px;
+	margin-top: 0;
+}
+
+p {
+	color: rgb(107, 114, 128);
+	font-size: 15px;
+	margin-bottom: 10px;
+	margin-top: 5px;
+}
+
+button {
+	font-size: 15px;
+}
+
+.softspan {
+	color: rgb(127, 134, 148);
+	font-size: 15px;
+	margin-bottom: 10px;
+	margin-top: 5px;
+}
+
+.card {
+	max-width: 800px;
+	margin: 0 auto;
+	padding: 16px;
+	border: 1px solid lightgray;
+	border-radius: 16px;
+}
+
+.card p:last-child {
+	margin-bottom: 0;
+}

style.css

@@ -0,0 +1,28 @@
+body {
+	padding: 2rem;
+	font-family: -apple-system, BlinkMacSystemFont, "Arial", sans-serif;
+}
+
+h1 {
+	font-size: 16px;
+	margin-top: 0;
+}
+
+p {
+	color: rgb(107, 114, 128);
+	font-size: 15px;
+	margin-bottom: 10px;
+	margin-top: 5px;
+}
+
+.card {
+	max-width: 620px;
+	margin: 0 auto;
+	padding: 16px;
+	border: 1px solid lightgray;
+	border-radius: 16px;
+}
+
+.card p:last-child {
+	margin-bottom: 0;
+}

templates/index.html

@@ -0,0 +1 @@
+../static/index.html

tokenizers_patch.py

@@ -0,0 +1,26 @@
+import os
+import sys
+import subprocess
+
+
+print("Getting rustup")
+subprocess.run(
+    "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y",
+    shell=True,
+)
+print("Got rustup")
+myenv = os.environ.copy()
+myenv["PATH"] = os.path.expanduser("~/.cargo/bin:") + myenv["PATH"]
+print("RUSTC", os.path.isfile(os.path.expanduser("~/.cargo/bin/rustc")))
+subprocess.run("rustc --version", shell=True, env=myenv)
+subprocess.run(
+    "pip install -e git+https://github.com/huggingface/tokenizers/#egg=tokenizers\&subdirectory=bindings/python",
+    shell=True,
+    env=myenv,
+)
+sys.path.append(
+    os.path.join(os.getcwd(), "src", "tokenizers", "bindings", "python", "py_src")
+)
+
+
+import tokenizers