fix query encoding and add new examples

modules/app.py

@@ -2,6 +2,7 @@ import sys
 from typing import List
 import traceback
 import os
+import base64
 # needs to be imported *before* transformers
 if os.path.exists('use_normal_tokenizers'):
     import tokenizers
@@ -11,8 +12,10 @@ else:
     import tokenizers_patch
     BIG_MODEL = True
     CUDA = True
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import json
+import pprint
 
 # from flask import Flask, request, render_template
 # from flask_cors import CORS
@@ -32,8 +35,14 @@ TRUNCATION_MESSAGE = f'warning: This demo is limited to {MAX_LENGTH} tokens in t
 
 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
@@ -43,7 +52,7 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 
 
 print("loading model")
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
 print("loading tokenizer")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 print("loading complete")
@@ -154,9 +163,18 @@ 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
+    print(info)
+    info = base64.urlsafe_b64decode(info + '=' * (4 - len(info) % 4)).decode('utf-8')
+    print(info)
     form = json.loads(info)
+    pprint.pprint(form)
+    # print(form)
     prompt = form['prompt']
     length_limit = int(form['length'])
     temperature = float(form['temperature'])
@@ -174,9 +192,17 @@ async def generate_maybe(info: str):
         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
+    print(info)
+    info = base64.urlsafe_b64decode(info + '=' * (4 - len(info) % 4)).decode('utf-8')
+    print(info)
     form = json.loads(info)
+    pprint.pprint(form)
     length_limit = int(form['length'])
     temperature = float(form['temperature'])
     max_retries = 1

static/index.html

@@ -134,6 +134,7 @@ label {
             <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");'>Multi-region</a></span>
             <span class="softspan"><a href='javascript:select_example("docstring");'>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>
@@ -252,12 +253,20 @@ def <infill>
         "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": 
-`<| file ext=.py |>
-
-def count_words(filename: str) -> <infill>
+`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 = {}
@@ -310,7 +319,7 @@ def count_words(filename):
         "mode": "python"
     },
     "javascript": {
-        "prompt": "<| file ext=.js |>\n // is something really happening here",
+        "prompt": "// fetch from the given URL and load the response contents into a new div",
         "length": 64,
         "temperature": 0.6,
         "mode": "javascript"
@@ -529,6 +538,7 @@ function make_generate_listener(url) {
             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"]);
@@ -540,6 +550,7 @@ function make_generate_listener(url) {
                     $("#warning").text("");
                 }
             } else {
+                console.log("error");
                 set_text(receive_data["text"])
                 $("#error").text(receive_data["message"]);
             }
@@ -552,13 +563,18 @@ function make_generate_listener(url) {
             $("#error").text(err);
         }
 
-        encoded_data = JSON.stringify(send_data)
+        encoded_data = encodeURIComponent(btoa(JSON.stringify(send_data)))
 
         try {
             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);
-                complete();
+                console.log("here");
+                console.log(response.status);
             } else {
                 response.json().then(success).catch(error).finally(complete);
             }