Upload 2 files

fn.py

@@ -86,7 +86,8 @@ def chatinterface_to_messages(message, history):
     messages = []
     
     if cfg['instruction']:
-        messages.append({'role': 'system', 'content': cfg['instruction']})
+        messages.append({'role': 'user', 'content': cfg['instruction']})
+        messages.append({'role': 'assistant', 'content': 'I understand.'})
 
     for pair in history:
         [user, assistant] = pair
@@ -100,32 +101,43 @@ def chatinterface_to_messages(message, history):
 
     return messages
 
-def apply_template(messages):
+def apply_template(message, history, args):
     global tokenizer, cfg
 
+    if 'input' in args:
+        message = args['input']
+    if 'instruction' in args:
+        cfg['instruction'] = args['instruction']
+
+    if 'messages' in args:
+        messages = args['messages']
+    elif history:
+        messages = chatinterface_to_messages(message, history)
+    else:
+        messages = {}
+
     if cfg['chat_template']:
         tokenizer.chat_template = cfg['chat_template']
 
-    if type(messages) is str:
+    if message:
         if cfg['inst_template']:
-            return cfg['inst_template'].format(instruction=cfg['instruction'], input=messages)
-        messages = [
-            {'role': 'user', 'content': cfg['instruction']},
-            {'role': 'assistant', 'content': 'I understand.'},
-            {'role': 'user', 'content': messages},
-        ]
-    if type(messages) is list:
-        return tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False)
-
-def chat(message, history = [], instruction = None, args = {}):
+            return cfg['inst_template'].format(instruction=cfg['instruction'], input=message)
+        if cfg['instruction']:
+            messages = [
+                {'role': 'user', 'content': cfg['instruction']},
+                {'role': 'assistant', 'content': 'I understand.'},
+                {'role': 'user', 'content': messages},
+            ]
+        else:
+            messages = [
+                {'role': 'user', 'content': message},
+            ]
+    return tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False)
+
+def chat(message = None, history = [], args = {}):
     global tokenizer, model, cfg
 
-    if instruction:
-        cfg['instruction'] = instruction
-        prompt = apply_template(message)
-    else:
-        messages = chatinterface_to_messages(message, history)
-        prompt = apply_template(messages)
+    prompt = apply_template(message, history, args)
 
     inputs = tokenizer(prompt, return_tensors="pt",
         padding=True, max_length=cfg['max_length'], truncation=True).to("cuda")
@@ -164,23 +176,40 @@ def chat(message, history = [], instruction = None, args = {}):
             # gradioは常に全文を返して欲しい
             yield model_output
 
-def infer(message, history = [], instruction = None, args = {}):
-    content = ''
-    for s in chat(message, history, instruction, args):
-        content += s
-    return content
+def infer(message = None, history = [], args = {}):
+    global tokenizer, model, cfg
+
+    prompt = apply_template(message, history, args)
+
+    inputs = tokenizer(prompt, return_tensors="pt",
+        padding=True, max_length=cfg['max_length'], truncation=True).to("cuda")
 
-def numel(message, history = [], instruction = None, args = {}):
+    generate_kwargs = dict(
+        inputs,
+        do_sample=True,
+        num_beams=1,
+        use_cache=True,
+    )
+
+    for k in [
+        'max_new_tokens',
+        'temperature',
+        'top_p',
+        'top_k',
+        'repetition_penalty'
+        ]:
+        if cfg[k]:
+            generate_kwargs[k] = cfg[k]
+
+    output_ids = model.generate(**generate_kwargs)
+    return tokenizer.decode(output_ids.tolist()[0][inputs['input_ids'].size(1):], skip_special_tokens=True)
+
+def numel(message = None, history = [], args = {}):
     global tokenizer, model, cfg
 
-    if instruction:
-        cfg['instruction'] = instruction
-        prompt = apply_template(message)
-    else:
-        messages = chatinterface_to_messages(message, history)
-        prompt = apply_template(messages)
+    prompt = apply_template(message, history, args)
 
-    model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+    model_inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     return torch.numel(model_inputs['input_ids'])
 

main.py

@@ -33,17 +33,17 @@ async def api_set_config(args: dict):
 
 @app.post("/infer")
 async def api_infer(args: dict):
-    args['fastapi'] = True
-    if 'stream' in args and args['stream']:
-        return StreamingResponse(
-            fn.chat(args['input'], [], args['instruct'], args),
-            media_type="text/event-stream",
-        )
-    else:
-        content = fn.infer(args['input'], [], args['instruct'], args)
-        return {'content': content}
+    content = fn.infer(args=args)
+    return {'content': content}
+
+@app.post("/stream")
+async def api_stream(args: dict):
+    return StreamingResponse(
+        fn.chat(args=args),
+        media_type="text/event-stream",
+    )
 
 @app.post("/numel")
 async def api_numel(args: dict):
-    content = fn.numel(args['input'], [], args['instruct'], args)
+    content = fn.numel(args=args)
     return {'numel': content}