[update]edit main

main.py

@@ -53,8 +53,6 @@ examples = [
 def main():
     args = get_args()
 
-    use_cpu = os.environ.get("USE_CPU", "all")
-
     tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path, trust_remote_code=True)
     # QWenTokenizer比较特殊, pad_token_id, bos_token_id, eos_token_id 均 为None. eod_id对应的token为<|endoftext|>
     if tokenizer.__class__.__name__ == "QWenTokenizer":
@@ -62,27 +60,21 @@ def main():
         tokenizer.bos_token_id = tokenizer.eod_id
         tokenizer.eos_token_id = tokenizer.eod_id
 
-    if not use_cpu:
-        model = AutoModel.from_pretrained(
-            args.pretrained_model_name_or_path,
-            trust_remote_code=True
-        ).half().cuda()
-    else:
-        model = AutoModelForCausalLM.from_pretrained(
-            args.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 = AutoModelForCausalLM.from_pretrained(
+        args.pretrained_model_name_or_path,
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        offload_folder="./offload",
+        offload_state_dict=True,
+        # load_in_4bit=True,
+    )
     model = model.eval()
 
-    def fn(inputs: str):
+    def fn(text: str):
         input_ids = tokenizer(
-            inputs,
+            text,
             return_tensors="pt",
             add_special_tokens=False,
         ).input_ids.to(args.device)
@@ -97,20 +89,28 @@ def main():
             response = response.strip().replace(tokenizer.eos_token, "").strip()
         return response
 
+    def fn_stream(text: str,
+                  max_new_tokens: int = 200,
+                  top_p: float = 0.85,
+                  temperature: float = 0.35,
+                  repetition_penalty: float = 1.2
+                  ):
+        return
+
     with gr.Blocks() as blocks:
         gr.Markdown(value=description)
 
         chatbot = gr.Chatbot([], elem_id="chatbot").style(height=400)
         with gr.Row():
             with gr.Column(scale=4):
-                text = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(container=False)
+                input_text_box = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(container=False)
             with gr.Column(scale=1):
                 button = gr.Button("Generate")
 
-        gr.Examples(examples, text)
+        gr.Examples(examples, input_text_box)
 
-        text.submit(fn, [text], [chatbot])
-        button.click(fn, [text], [chatbot])
+        input_text_box.submit(fn, [input_text_box], [chatbot])
+        button.click(fn, [input_text_box], [chatbot])
 
     blocks.queue().launch()