[update]add main

main.py

@@ -16,39 +16,63 @@ def greet(question: str, history: List[Tuple[str, str]]):
     return result
 
 
+model_map: dict = dict()
+
+
+def init_model(pretrained_model_name_or_path: str, device: str):
+    global model_map
+    if pretrained_model_name_or_path not in model_map.keys():
+        # clear
+        for k1, v1 in model_map.items():
+            for k2, v2 in v1.items():
+                del v2
+        model_map = dict()
+
+        # build model
+        model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            offload_folder="./offload",
+            offload_state_dict=True,
+            # load_in_4bit=True,
+        )
+        model = model.to(device)
+        model = model.bfloat16().eval()
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            trust_remote_code=True,
+            # llama不支持fast
+            use_fast=False if model.config.model_type == "llama" else True,
+            padding_side="left"
+        )
+
+        # QWenTokenizer比较特殊, pad_token_id, bos_token_id, eos_token_id 均 为None. eod_id对应的token为<|endoftext|>
+        if tokenizer.__class__.__name__ == "QWenTokenizer":
+            tokenizer.pad_token_id = tokenizer.eod_id
+            tokenizer.bos_token_id = tokenizer.eod_id
+            tokenizer.eos_token_id = tokenizer.eod_id
+
+        model_map[pretrained_model_name_or_path] = {
+            "model": model,
+            "tokenizer": tokenizer,
+        }
+    else:
+        model = model_map[pretrained_model_name_or_path]["model"]
+        tokenizer = model_map[pretrained_model_name_or_path]["tokenizer"]
+    return model, tokenizer
+
+
 def chat_with_llm_non_stream(question: str,
                              history: List[Tuple[str, str]],
                              pretrained_model_name_or_path: str,
                              max_new_tokens: int, top_p: float, temperature: float, repetition_penalty: float,
+                             device: str
                              ):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    model = AutoModelForCausalLM.from_pretrained(
-        pretrained_model_name_or_path,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        offload_folder="./offload",
-        offload_state_dict=True,
-        # load_in_4bit=True,
-    )
-    model = model.to(device)
-    model = model.bfloat16().eval()
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        pretrained_model_name_or_path,
-        trust_remote_code=True,
-        # llama不支持fast
-        use_fast=False if model.config.model_type == "llama" else True,
-        padding_side="left"
-    )
-
-    # QWenTokenizer比较特殊, pad_token_id, bos_token_id, eos_token_id 均 为None. eod_id对应的token为<|endoftext|>
-    if tokenizer.__class__.__name__ == "QWenTokenizer":
-        tokenizer.pad_token_id = tokenizer.eod_id
-        tokenizer.bos_token_id = tokenizer.eod_id
-        tokenizer.eos_token_id = tokenizer.eod_id
+    model, tokenizer = init_model(pretrained_model_name_or_path, device)
 
     input_ids = tokenizer(
         question,
@@ -70,10 +94,11 @@ def chat_with_llm_non_stream(question: str,
             eos_token_id=tokenizer.eos_token_id
         )
         outputs = outputs.tolist()[0][len(input_ids[0]):]
-        response = tokenizer.decode(outputs)
-        response = response.strip().replace(tokenizer.eos_token, "").strip()
+        answer = tokenizer.decode(outputs)
+        answer = answer.strip().replace(tokenizer.eos_token, "").strip()
 
-    return
+    result = history + [(question, answer)]
+    return result
 
 
 def main():
@@ -81,8 +106,10 @@ def main():
     chat llm
     """
 
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+
     with gr.Blocks() as blocks:
-        gr.Markdown(value="gradio demo")
+        gr.Markdown(value=description)
 
         chatbot = gr.Chatbot([], elem_id="chatbot", height=400)
         with gr.Row():
@@ -115,7 +142,8 @@ def main():
 
         inputs = [
             text_box, chatbot, model_name,
-            max_new_tokens, top_p, temperature, repetition_penalty
+            max_new_tokens, top_p, temperature, repetition_penalty,
+            device
         ]
         outputs = [
             chatbot