Update app.py

app.py

@@ -53,10 +53,10 @@ class PromptEnhancerV2:
         self,
         prompt_cot,
         sys_prompt="请根据用户的输入，生成思考过程的思维链并改写提示词：",
-        temperature=0.0,
+        temperature=0.1,
         top_p=1.0,
         max_new_tokens=2048,
-        device="cuda:0",
+        device="cuda",
     ):
         org_prompt_cot = prompt_cot
         try:
@@ -112,91 +112,6 @@ class PromptEnhancerV2:
             print(f"✗ Re-prompting failed, so we are using the original prompt. Error: {e}")
 
         return prompt_cot
-    # @torch.inference_mode()
-    @spaces.GPU
-    def predict_stream(
-        self,
-        prompt_cot,
-        sys_prompt="请根据用户的输入，生成思考过程的思维链并改写提示词：",
-        temperature=0.1,
-        top_p=1.0,
-        max_new_tokens=2048,
-        device="cuda:0",
-    ):
-        org_prompt_cot = prompt_cot
-
-        # 组装输入，同 predict
-        user_prompt_format = sys_prompt + "\n" + org_prompt_cot
-        messages = [{"role": "user", "content": [{"type": "text", "text": user_prompt_format}]}]
-        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        image_inputs, video_inputs = process_vision_info(messages)
-        inputs = self.processor(
-            text=[text],
-            images=image_inputs,
-            videos=video_inputs,
-            padding=True,
-            return_tensors="pt",
-        )
-        inputs = inputs.to(device)
-
-        # 取得 tokenizer（大多数情况下 processor.tokenizer 就有；加一个后备以防万一）
-        tokenizer = getattr(self.processor, "tokenizer", None)
-        if tokenizer is None:
-            tokenizer = AutoTokenizer.from_pretrained(self.models_root_path, trust_remote_code=True)
-
-        streamer = TextIteratorStreamer(
-            tokenizer=tokenizer,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False,
-        )
-
-        gen_kwargs = dict(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=float(temperature),
-            do_sample=True,   # 与原逻辑一致; 若要采样流式把这里改为 True
-            top_k=5,
-            top_p=0.9,
-            streamer=streamer,
-        )
-
-        # 子线程启动生成；主线程消费 streamer
-        thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
-        thread.start()
-
-        buffer = ""     # 累积完整输出（含思考）
-        emitted = ""    # 已对外输出的“重写提示词”部分
-        already_stripped_newline = False
-
-        try:
-            for piece in streamer:
-                buffer += piece
-                part = buffer.split('assistant')[-1]
-                delta = part[len(emitted):]
-                if delta:
-                    emitted = part
-                    yield emitted  # 将中间结果送给前端
-        finally:
-            thread.join()
-
-        # 如果始终没等到第二个 think>，回退到原始 prompt
-        # if emitted.strip() == "":
-        #     yield replace_single_quotes(org_prompt_cot)
-        try:
-            assert emitted.count("think>") == 2
-            prompt_cot = emitted.split("think>")[-1]
-            if prompt_cot.startswith("\n"):
-                prompt_cot = prompt_cot[1:]
-            prompt_cot = emitted.split('assistant')[-1] + '\n \n Recaption:'+replace_single_quotes(prompt_cot)
-            # prompt_cot = replace_single_quotes(prompt_cot)
-            yield prompt_cot
-        except Exception as e:
-            prompt_cot = org_prompt_cot
-            print(f"✗ Re-prompting failed, so we are using the original prompt. Error: {e}")
-            yield prompt_cot
-
-
-
 # -------------------------
 # Gradio app helpers
 # -------------------------
@@ -223,32 +138,6 @@ def ensure_enhancer(state, model_path, device_map, torch_dtype):
         return {"enhancer": enhancer, "model_path": model_path, "device_map": device_map, "torch_dtype": torch_dtype}
     return state
 
-def stream_single(prompt, sys_prompt, temperature, max_new_tokens, device,
-                  model_path, device_map, torch_dtype, state):
-    if not prompt or not str(prompt).strip():
-        yield "", "请先输入提示词。", state
-        return
-
-    t0 = time.time()
-    state = ensure_enhancer(state, model_path, device_map, torch_dtype)
-    enhancer = state["enhancer"]
-
-    emitted = ""
-    try:
-        for chunk in enhancer.predict_stream(
-            prompt_cot=prompt,
-            sys_prompt=sys_prompt,
-            temperature=temperature,
-            max_new_tokens=max_new_tokens,
-            device=device
-        ):
-            emitted = chunk
-            info = f"已接收 {len(emitted)} 字符，用时 {time.time()-t0:.2f}s"
-            yield emitted, info, state
-        # 结束时再给一次最终状态（可选）
-        yield emitted, f"完成。总耗时 {time.time()-t0:.2f}s", state
-    except Exception as e:
-        yield "", f"推理失败：{e}", state
 
 def run_single(prompt, sys_prompt, temperature, max_new_tokens, device,
                model_path, device_map, torch_dtype, state):
@@ -295,11 +184,12 @@ with gr.Blocks(title="Prompt Enhancer_V2") as demo:
                 value=DEFAULT_MODEL_PATH,
                 placeholder="/apdcephfs_jn3/share_302243908/aladdinwang/model_weight/cot_taurus_v6_50/global_step0",
             )
-            device_map = gr.Dropdown(
-                choices=["auto", "cuda", "cpu"],
-                value="auto",
-                label="device_map（模型加载映射）"
-            )
+            device_map = "cuda"
+            # gr.Dropdown(
+            #     choices=["auto", "cuda", "cpu"],
+            #     value="auto",
+            #     label="device_map（模型加载映射）"
+            # )
             torch_dtype = gr.Dropdown(
                 choices=["bfloat16", "float16", "float32"],
                 value="bfloat16",