Upload 3 files

app.py

@@ -96,10 +96,16 @@ with gr.Blocks() as demo:
             with gr.Column(scale=1):
                 said = gr.Textbox(
                     label='said',
-                    lines=25,
+                    lines=20,
+                    show_copy_button=True,
+                    )
+                numel = gr.Textbox(
+                    lines=1,
+                    label='numel',
                     show_copy_button=True,
                     )
         inst_button = gr.Button(value='inst')
+        numel_button = gr.Button(value='numel')
 
     with gr.Tab('chat'):
         gr.ChatInterface(fn.chat)
@@ -116,5 +122,11 @@ with gr.Blocks() as demo:
         outputs=[said],
         )
 
+    numel_button.click(
+        fn=fn.numel,
+        inputs=[input, input, instruction],
+        outputs=[numel],
+        )
+
 if __name__ == '__main__':
     demo.launch()

fn.py

@@ -180,6 +180,9 @@ def chat(message, history = [], instruction = None, args = {}):
         prompt = apply_template(messages)
 
     model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+    # どうせ0固定だしエラーが出るので消してしまう
+    if 'token_type_ids' in model_inputs:
+        del model_inputs['token_type_ids']
 
     streamer = TextIteratorStreamer(
         tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True,
@@ -220,3 +223,17 @@ def infer(message, history = [], instruction = None, args = {}):
     for s in chat(message, history, instruction, args):
         content += s
     return content
+
+def numel(message, history = [], instruction = None, 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)
+
+    model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+
+    return torch.numel(model_inputs['input_ids'])

main.py

@@ -42,3 +42,8 @@ async def api_infer(args: dict):
     else:
         content = fn.infer(args['input'], [], args['instruct'], args)
         return {'content': content}
+
+@app.post("/numel")
+async def api_numel(args: dict):
+    content = fn.numel(args['input'], [], args['instruct'], args)
+    return {'numel': content}