add stream support

Author: little51

chat/package.json

@@ -1,6 +1,6 @@
 {
   "name": "chat",
-  "version": "0.1.0",
+  "version": "1.0.0",
   "private": true,
   "dependencies": {
     "@chatui/core": "^2.4.2",

chat/src/App.js

@@ -1,7 +1,9 @@
 import './App.css';
-import Chat, { Bubble, useMessages } from '@chatui/core';
+import Chat, { Bubble, useMessages, Progress } from '@chatui/core';
 import '@chatui/core/dist/index.css';
-import React, { useEffect } from 'react'
+import '@chatui/core/es/styles/index.less';
+import React, { useEffect, useState } from 'react';
+import './chatui-theme.css';
 
 const defaultQuickReplies = [
   {
@@ -48,19 +50,20 @@ const initialMessages = [
 
 function App() {
   const { messages, appendMsg, setTyping } = useMessages(initialMessages);
+  const [percentage, setPercentage] = useState(0);
 
   function handleSend(type, val) {
     if (type === 'text' && val.trim()) {
       appendMsg({
         type: 'text',
         content: { text: val },
-        position: 'right',
+        position: 'left',
         user: { avatar: '//gitclone.com/download1/user.png' },
       });
 
       setTyping(true);
-
-      onGenCode(val);
+      setPercentage(10);
+      onGenCode(val, val, 0);
     }
   }
 
@@ -90,7 +93,7 @@ function App() {
     } else if (item.name === "Java") {
       content = "int add(int x,int y){";
     } else if (item.name === "javascript") {
-      content = "function Add(x,y,z){";
+      content = "function Add(x,y){";
     } else if (item.name === "golang") {
       content = "func IsBlacklist(bl []string,url string) bool{";
     } else {
@@ -99,26 +102,57 @@ function App() {
     handleSend('text', content);
   }
 
-  function onGenCode(context) {
-    var sl = context.trim().split("\n");
-    context = sl[sl.length - 1];
-    if (context.trim() === "") {
-      alert("输入不能为空！")
+  function onGenCode(context_en, context_ch, count) {
+    if (count >= 5) {
+      setPercentage(0);
       return;
     }
     let xhr = new XMLHttpRequest();
-    xhr.open('post', 'https://gitclone.com/aiit/codegen');
+    xhr.open('post', 'https://gitclone.com/aiit/codegen_stream');
     xhr.setRequestHeader('Content-Type', 'application/json');
     xhr.onload = function () {
       var json = JSON.parse(xhr.response);
-      context = context + "\n" + json.result;
-      appendMsg({
-        type: 'text',
-        content: { text: context },
-        user: { avatar: '//gitclone.com/download1/gitclone.png' },
-      });
+      if (count === 0) {
+        context_en = context_en + "\n" + json.result_en;
+        context_ch = context_ch + "\n" + json.result_ch;
+        appendMsg({
+          type: 'text',
+          content: { text: context_ch },
+          user: { avatar: '//gitclone.com/download1/gitclone.png' },
+        });
+      } else {
+        if (("" === json.result_en.trim()) || json.result_en.trim().startsWith("A:") || json.result_en.trim().endsWith("A:")) {
+          setPercentage(0);
+          return;
+        }
+        context_en = context_en + json.result_en;
+        context_ch = context_ch + json.result_ch;
+        if (context_ch === context_en) {
+          updateMsg(context_en);
+        } else {
+          updateMsg(context_ch + "\n" + context_en);
+        }
+
+      }
+      count++;
+      setPercentage(count * 20);
+      onGenCode(context_en, context_ch, count);
+    }
+    xhr.send(JSON.stringify({
+      "context": context_en,
+      "maxlength": 16,
+      "modelname": "codegen"
+    }));
+
+    function updateMsg(context_ch) {
+      var oUl = document.getElementById('root');
+      var aBox = getByClass(oUl, 'Bubble text');
+      if (aBox.length > 0) {
+        aBox[aBox.length - 1].innerHTML = "<p>" + context_ch + "</p>";
+        var msgList = getByClass(oUl, "PullToRefresh")[0];
+        msgList.scrollTo(0, msgList.scrollHeight);
+      }
     }
-    xhr.send('{"context":"' + context + '","maxlength":32}');
   }
 
   function findInArr(arr, n) {
@@ -169,14 +203,15 @@ function App() {
               title: 'More',
             },
           ],
-          title: '基于Salesforce codegen和GPTJ、GPT-neo的AI代码生成',
+          title: '基于Salesforce codegen和GPTJ的AI代码生成',
         }}
         messages={messages}
         renderMessageContent={renderMessageContent}
         quickReplies={defaultQuickReplies}
         onQuickReplyClick={handleQuickReplyClick}
         onSend={handleSend}
       />
+      <Progress value={percentage} />
     </div>
   );
 }

chat/src/chatui-theme.css

@@ -0,0 +1,13 @@
+:root {
+  font-size: 16px;
+  line-height:14px ;
+}
+.ChatApp,
+.MessageContainer,
+.Navbar,
+.Message .Bubble,
+.QuickReplies,
+.ChatFooter {
+  background-repeat: no-repeat;
+  background-size: cover;
+}

codegen.py

@@ -7,7 +7,8 @@
 from functools import lru_cache
 from aiohttp import web
 from jaxformer.hf.sample import load_model,sampling
-from gpt_neo import gpt_load_model,gpt_generate
+from gpt_j import gpt_load_model,gpt_generate
+from codegen_stream import codegen_stream
 
 ROOT = os.path.dirname(__file__)
 
@@ -18,24 +19,9 @@ async def index(request):
     return web.Response(content_type="text/html", text=content)
 
 @lru_cache(maxsize=1024, typed=False)
-def getAnswerFromChatGPT(context):
-    url = 'http://chatgptserver.com:5000/chat'
-    data = '{"message":"' +  context + '", "user": "gitclone"}'
-    headers = {'content-type': 'application/json;charset=utf-8'}
-    r = requests.post(url,data= data.encode(), headers=headers)
-    res = r.json()
-    return res['response']
-
-@lru_cache(maxsize=1024, typed=False)
-def getAnswerFromChatGPTJ(context):
-    #url = 'http://52.82.67.116:8081/generate/'
-    #data = '{' + '"text": "' + context + '",' + '"generate_tokens_limit": 40,'+ '"top_p": 0.7,'+'"top_k": 0,' + '"temperature":1.0' +'}' ;
-    #headers = {'content-type': 'application/json;charset=utf-8'}
-    #r = requests.post(url,data= data.encode(), headers=headers)
-    #res = r.json()
-    #return res['completion']
+def getAnswerFromChatGPTJ(context,maxlength):
     gpt_load_model()
-    return gpt_generate(context,128)
+    return gpt_generate(context,maxlength)
 
 async def codegen(request):
     params = await request.json()
@@ -56,7 +42,7 @@ async def codegen(request):
     print(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime()),"context : " + context)
     context = context.replace("//","").replace("#","").strip()
     if flag_chs :#or content.startwith('gpt-j') :
-        result = getAnswerFromChatGPTJ(context).replace(context,"")
+        result = getAnswerFromChatGPTJ(context,maxlength).replace(context,"")
     else:
         result = sampling(context,maxlength)
     end = time.perf_counter()
@@ -73,6 +59,7 @@ async def codegen(request):
 app.router.add_get("/", index)
 app.router.add_get("/codegen", index)
 app.router.add_post("/codegen", codegen)
+app.router.add_post("/codegen_stream", codegen_stream)
 
 for route in list(app.router.routes()):
     cors.add(route, {

codegen_stream.py

@@ -0,0 +1,42 @@
+import requests
+import time
+from aiohttp import web
+import json
+from jaxformer.hf.sample import load_model, sampling
+from gpt_j import gpt_load_model, gpt_generate_stream
+
+
+def sampling_gptj(context, maxlength):
+    gpt_load_model()
+    return gpt_generate_stream(context, maxlength)
+
+
+async def codegen_stream(request):
+    params = await request.json()
+    context = params["context"]
+    maxlength = params["maxlength"]
+    modelname = params["modelname"]
+    start = time.perf_counter()
+    print(time.strftime("%Y-%m-%d %H:%M:%S",
+          time.localtime()), "context : " + context)
+    context = context.strip()
+    f = lambda x='ddd': sum(
+        [1 if u'\u4e00' <= i <= u'\u9fff' else 0 for i in x]) > 0
+    flag_chs = f(context)
+    if flag_chs:
+        results = sampling_gptj(context, maxlength)
+        results = json.loads(results)
+        result_en = results["result_en"]
+        result_ch = results["result_ch"]
+    else:
+        result_en = sampling(context, maxlength)
+        result_ch = result_en
+    end = time.perf_counter()
+    print(time.strftime("%Y-%m-%d %H:%M:%S",
+          time.localtime()), "result  : " + result_ch)
+    return web.Response(
+        content_type="application/json",
+        text=json.dumps(
+            {"result_en": result_en, "result_ch": result_ch, "time": end-start}
+        ),
+    )

gpt_neo.py gpt_j.pygpt_neo.py renamed to gpt_j.py

@@ -3,17 +3,18 @@
 import gradio as gr
 import torch
 import requests
+import json
 
-#generator = None
+# generator = None
 translator_zh2en = None
 translator_en2zh = None
 
 
 def gpt_load_model():
-    #global generator
+    # global generator
     global translator_zh2en
     global translator_en2zh
-    #if generator is None:
+    # if generator is None:
     #    #torch.cuda.set_device('cuda:1')
     #    generator = pipeline(
     #        'text-generation', model='EleutherAI/gpt-neo-1.3B')
@@ -24,17 +25,21 @@ def gpt_load_model():
         translator_en2zh = pipeline(
             "translation", model="Helsinki-NLP/opus-mt-en-zh")
 
-def getAnswerFromChatGPTJ6B(context):
-    url = 'http://127.0.0.1:8081/generate/'
-    data = '{' + '"text": "' + context + '",' + '"generate_tokens_limit": 64,'+ '"top_p": 0.7,'+'"top_k": 0,' + '"temperature":0.9' +'}' ;
+
+def getAnswerFromChatGPTJ6B(context, maxlength):
+    url = 'http://172.16.62.66:8081/generate/'
+    data = '{' + '"text": "' + context + '",' + '"generate_tokens_limit": ' + \
+        str(maxlength) + ',' + '"top_p": 0.7,' + \
+        '"top_k": 0,' + '"temperature":0.9' + '}'
     headers = {'content-type': 'application/json;charset=utf-8'}
-    r = requests.post(url,data= data.encode(), headers=headers)
+    r = requests.post(url, data=data.encode(), headers=headers)
     res = r.json()
     return res['completion']
 
+
 @lru_cache(maxsize=1024, typed=False)
 def gpt_generate(inputs, maxlength):
-    #global generator
+    # global generator
     global translator_zh2en
     global translator_en2zh
     f = lambda x='ddd': sum(
@@ -44,25 +49,48 @@ def gpt_generate(inputs, maxlength):
     if flag_chs:
         inputs = translator_zh2en(inputs)[0]['translation_text']
         print("zh2en: ", inputs)
-    #results = generator(inputs, max_length=int(maxlength),
-    #                    do_sample=True, temperature=0.9)
-    results = getAnswerFromChatGPTJ6B(inputs)
+    results = getAnswerFromChatGPTJ6B(inputs, maxlength)
     print("output: ", results)
     if flag_chs:
-        #results = translator_en2zh(results[0]['generated_text'])
         results_en = results
         results = translator_en2zh(results)
         print("en2zh:", results)
         return results_en + '\n' + results[0]['translation_text']
     else:
         return results
 
+
+def gpt_generate_stream(inputs, maxlength):
+    # global generator
+    global translator_zh2en
+    global translator_en2zh
+    f = lambda x='ddd': sum(
+        [1 if u'\u4e00' <= i <= u'\u9fff' else 0 for i in x]) > 0
+    print("inputs: ", inputs)
+    flag_chs = f(inputs)
+    if flag_chs:
+        inputs = translator_zh2en(inputs)[0]['translation_text']
+        print("zh2en: ", inputs)
+    results = getAnswerFromChatGPTJ6B(inputs,maxlength)
+    print("output: ", results)
+    if flag_chs:
+        results_en = results
+        results = translator_en2zh(results)
+        print("en2zh:", results)
+        return json.dumps(
+            {"result_en": results_en, "result_ch":  results[0]['translation_text']})
+    else:
+        return json.dumps(
+            {"result_en": results, "result_ch":  results})
+
+
 def chat(message, history):
     history = history or []
-    response = gpt_generate(message,128)
+    response = gpt_generate(message, 128)
     history.append((message, response))
     return history, history
 
+
 def create_ui():
     chatbot = gr.Chatbot().style(color_map=("green", "gray"))
     interface = gr.Interface(
@@ -73,6 +101,7 @@ def create_ui():
     )
     interface.launch(server_name='0.0.0.0')
 
+
 if __name__ == "__main__":
     torch.cuda.set_device(1)
     print("torch gpu: ", torch.cuda.is_available())

jaxformer/hf/sample.py

@@ -67,7 +67,7 @@ def cast(model, fp16=True):
 
 def create_model(ckpt, fp16=True):
     if fp16:
-        return CodeGenForCausalLM.from_pretrained(ckpt, revision='float16', torch_dtype=torch.float16, low_cpu_mem_usage=True)
+        return CodeGenForCausalLM.from_pretrained(ckpt, revision='float16', torch_dtype=torch.float16, low_cpu_mem_usage=False)
     else:
         return CodeGenForCausalLM.from_pretrained(ckpt)