Update fn.py

fn.py

@@ -168,16 +168,20 @@ def chat(message, history = [], instruction = None, args = {}):
 
     model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
 
-    streamer = TextIteratorStreamer(
-        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True,
-    )
+    if 'fastapi' not in args or 'stream' in args and args['stream']:
+        streamer = TextIteratorStreamer(
+            tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True,
+        )
 
     generate_kwargs = dict(
         model_inputs,
-        streamer=streamer,
         do_sample=True,
-        num_beams=1,
     )
+
+    if 'fastapi' not in args or 'stream' in args and args['stream']:
+        generate_kwargs['streamer'] = streamer
+        generate_kwargs['num_beams'] = 1
+
     for k in [
         'max_new_tokens',
         'temperature',
@@ -188,43 +192,21 @@ def chat(message, history = [], instruction = None, args = {}):
         if cfg[k]:
             generate_kwargs[k] = cfg[k]
 
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    model_output = ""
-    for new_text in streamer:
-        model_output += new_text
-        if 'fastapi' in args:
-            # fastapiは差分だけを返して欲しい
-            yield new_text
-        else:
-            # gradioは常に全文を返して欲しい
-            yield model_output
-
-    return model_output
-
-def infer(args: dict):
-    global cfg
-
-    if 'model_name' in args:
-        load_model(args['model_name'], args['qtype'], args['dtype'])
+    if 'fastapi' not in args or 'stream' in args and args['stream']:
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
 
-    for k in [
-        'instruction',
-        'inst_template',
-        'chat_template',
-        'max_new_tokens',
-        'temperature',
-        'top_p',
-        'top_k',
-        'repetition_penalty'
-        ]:
-        cfg[k] = args[k]
+        model_output = ""
+        for new_text in streamer:
+            model_output += new_text
+            if 'fastapi' in args:
+                # fastapiは差分だけを返して欲しい
+                yield new_text
+            else:
+                # gradioは常に全文を返して欲しい
+                yield model_output
 
-    if 'messages' in args:
-        return chat(args['input'], args['messages'])
-    if 'instruction' in args:
-        return instruct(args['instruction'], args['input'])
+    return model.generate(**generate_kwargs)
 
 def apply_template(messages):
     global tokenizer, cfg
@@ -235,6 +217,6 @@ def apply_template(messages):
     if type(messages) is str:
         if cfg['inst_template']:
             return cfg['inst_template'].format(instruction=cfg['instruction'], input=messages)
-        return cfg['instruction']
+        return cfg['instruction'].format(input=messages)
     if type(messages) is list:
         return tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False)