Update README.md

README.md

@@ -77,200 +77,11 @@ from peft import LoraConfig, get_peft_model, PeftModel
 
 from qwen_omni_utils import process_mm_info
 
-# ============================================================
-# Main Agent Stages
-# ============================================================
-def stage1_collector(model, processor, video_path, query, device):
-    """Collector uses base_model.thinker with LoRA disabled"""
-    has_audio = check_if_video_has_audio(video_path)
-    content = [{"type": "video", "video": video_path, 
-                "min_pixels": 36 * 28 * 28,
-                "max_pixels": 96 * 28 * 28,
-                "max_frames": 150, 
-                "fps": 1.0}]
-    if has_audio:
-        content.append({"type": "audio", "audio": video_path})
-    text_prompt = COLLECTOR_USER_TEMPLATE.format(problem=query)
-    content.append({"type": "text", "text": text_prompt})
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": COLLECTOR_SYSTEM_PROMPT}]},
-        {"role": "user", "content": content}
-    ]
-    print(f"collector input: {messages}")
-
-    audios, images, videos = process_mm_info(messages, use_audio_in_video=False)
-    text = processor.apply_chat_template([messages], tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True).to(device)
-
-    for k in inputs:
-        if isinstance(inputs[k], torch.Tensor) and inputs[k].is_floating_point():
-            inputs[k] = inputs[k].to(torch.bfloat16)
-
-    with torch.inference_mode():
-        output_ids = model.generate(**inputs, max_new_tokens=512, do_sample=False)
-    resp = processor.decode(output_ids[0][inputs.input_ids.size(1):], skip_special_tokens=True)
-    return resp
-
-def stage2_planner(model, processor, video_path, query, collector_text, device):
-    """Planner uses base_model.thinker with LoRA enabled"""
-    has_audio = check_if_video_has_audio(video_path)
-    content = [{"type": "video", 
-                "video": video_path, 
-                "min_pixels": 36 * 28 * 28, 
-                "max_pixels": 64 * 28 * 28,
-                "max_frames": 150, 
-                "fps": 1.0}]
-    if has_audio:
-        content.append({"type": "audio", "audio": video_path})
-    text_prompt = PLANNER_USER_TEMPLATE.format(collector_text=collector_text, problem=query)
-    content.append({"type": "text", "text": text_prompt})
-
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": PLANNER_SYSTEM_PROMPT}]},
-        {"role": "user", "content": content}
-    ]
-    print(f"planner input: {messages}")
-
-    audios, images, videos = process_mm_info(messages, use_audio_in_video=False)
-    text = processor.apply_chat_template([messages], tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True).to(device)
-
-    for k in inputs:
-        if isinstance(inputs[k], torch.Tensor) and inputs[k].is_floating_point():
-            inputs[k] = inputs[k].to(torch.bfloat16)
-
-    with torch.inference_mode():
-        output_ids = model.generate(**inputs, max_new_tokens=256, do_sample=False)
-    resp = processor.decode(output_ids[0][inputs.input_ids.size(1):], skip_special_tokens=True)
-    return parse_planner_output(resp), resp
-
-def stage3_grounder(grounder_proc, video_path, query, duration):
-    req = json.dumps({"video_path": video_path, "query": query, "duration": duration})
-    try:
-        grounder_proc.stdin.write(req + "\n")
-        grounder_proc.stdin.flush()
-    except (BrokenPipeError, OSError) as exc:
-        raise RuntimeError(f"Grounder worker unavailable while sending request: {exc}") from exc
-
-    resp_line = grounder_proc.stdout.readline().strip()
-    if not resp_line:
-        raise RuntimeError("Grounder worker returned empty response.")
-
-    try:
-        resp = json.loads(resp_line)
-    except json.JSONDecodeError as exc:
-        raise RuntimeError(f"Grounder worker returned invalid JSON: {resp_line}") from exc
-
-    if "error" in resp and resp["error"]:
-        raise RuntimeError(f"Grounder worker error: {resp['error']}")
-
-    return resp.get("pred_spans", [[0, duration]]), resp.get("success", False)
-
-def build_humanomni_query(sample):
-    """Build the full question text with options and TYPE_TEMPLATE, aligned with eval_humanomniv2.py."""
-    problem_type = sample.get("problem_type", "free-form")
-    if problem_type in ("multiple choice", "emer_ov_mc"):
-        question = sample["problem"] + " Options:\n"
-        for op in sample.get("options", []):
-            question += op + "\n"
-    else:
-        question = sample["problem"]
-    question += TYPE_TEMPLATE.get(problem_type, TYPE_TEMPLATE["free-form"])
-    print(f"[HumanOmni question] {question}")
-    return question
-
-def get_humanomni_inputs(processor, video_path, query_text, sample, device):
-    """Build model inputs for HumanOmniV2, aligned with eval_humanomniv2.py message format."""
-    has_audio = check_if_video_has_audio(video_path)
-    data_type = sample.get("data_type", "video")
-    content = [{"type": data_type, 
-                data_type: video_path,
-                "min_pixels": 128 * 28 * 28,
-                "max_pixels": 256 * 28 * 28,
-                "max_frames": 64,
-                "fps": 2.0}]
-    if has_audio:
-        content.append({"type": "audio", "audio": video_path})
-        text_prompt = f"Here is a {data_type}, with the audio from the video.\n" + query_text
-    else:
-        text_prompt = query_text
-    content.append({"type": "text", "text": text_prompt})
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": HUMANOMNI_SYSTEM_PROMPT}]},
-        {"role": "user", "content": content}
-    ]
-    print(f"humanomni input: {messages}")
-    audios, images, videos = process_mm_info(messages, use_audio_in_video=False)
-    text = processor.apply_chat_template([messages], tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True).to(device)
-
-    for k in inputs:
-        if isinstance(inputs[k], torch.Tensor) and inputs[k].is_floating_point():
-            inputs[k] = inputs[k].to(torch.bfloat16)
-    return inputs
-
-def revise_answer(model, processor, video_path, query, answer, device):
-    """Reviser uses base_model.thinker with LoRA disabled"""
-    has_audio = check_if_video_has_audio(video_path)
-    content = [{"type": "video", 
-                "video": video_path,
-                "min_pixels": 64 * 28 * 28,
-                "max_pixels": 128 * 28 * 28,
-                "max_frames": 100,
-                "fps": 1.0}]
-    if has_audio:
-        content.append({"type": "audio", "audio": video_path})
-    user_prompt = f"Query: {query}\nCandidate Answer: {answer}"
-    content.append({"type": "text", "text": user_prompt})
-
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": REVISER_SYSTEM_PROMPT}]},
-        {"role": "user", "content": content}
-    ]
-
-    audios, images, videos = process_mm_info(messages, use_audio_in_video=False)
-    text = processor.apply_chat_template([messages], tokenize=False, add_generation_prompt=True)
-    inputs = processor(text=text, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True).to(device)
-
-    for k in inputs:
-        if isinstance(inputs[k], torch.Tensor) and inputs[k].is_floating_point():
-            inputs[k] = inputs[k].to(torch.bfloat16)
-
-    with torch.inference_mode():
-        output_ids = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
-    resp = processor.decode(output_ids[0][inputs.input_ids.size(1):], skip_special_tokens=True)
-    return extract_score(resp), resp
-
 # ============================================================
 # Main Process
 # ============================================================
 def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--base_model_path", required=True)
-    parser.add_argument("--planner_lora_path", required=True)
-    parser.add_argument("--humanomni_path", required=True)
-    parser.add_argument("--grounder_path", required=True)
-    parser.add_argument("--grpo_adapter_path", required=True)
-    parser.add_argument("--grounder_python", required=True)
-    parser.add_argument("--grounder_cwd", required=True)
-    parser.add_argument("--input_file", required=True)
-    parser.add_argument("--output_file", required=True)
-    parser.add_argument("--lora_save_dir", required=True)
-    parser.add_argument("--video_root", default="", help="Root directory for video files")
-    parser.add_argument("--id_key", default="qid,video", help="Comma separated keys for identifying samples")
-    parser.add_argument("--main_gpu", default="cuda:0")
-    parser.add_argument("--grounder_gpu", default="cuda:1")
-    parser.add_argument("--humanomni_gpu", default="cuda:0")
-    parser.add_argument("--t_max", type=int, default=3)
-    parser.add_argument("--tau", type=int, default=7)
-    parser.add_argument("--alpha", type=float, default=0.7)
-    parser.add_argument("--b0", type=float, default=5.0)
-    parser.add_argument("--lr", type=float, default=1e-5)
-    parser.add_argument("--num_chunks", type=int, default=1, help="Total number of parallel processes")
-    parser.add_argument("--chunk_idx", type=int, default=0, help="Current process index (0-based)")
-    args = parser.parse_args()
    
-
     # ---- Initialize Models ----
     print(f"\n[INIT] Loading Base Model ({args.base_model_path}) on {args.main_gpu}")
     base_model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
@@ -322,16 +133,6 @@ def main():
 
     # ---- 3. Loop through dataset ----
     for sample in samples_to_process:
-        dataset_id = get_sample_id(sample, args.id_key)
-        optimizer = None
-        query = sample.get("problem")
-        video_path = sample.get("video") or sample.get("path")
-        if args.video_root and video_path:
-            video_path = os.path.join(args.video_root, video_path)
-        trim_path = None
-        adapter_name = None
-        adapter_deleted = False
-
         try:
             print(f"\n======== Processing [{dataset_id}] ========")
             duration = get_video_duration(video_path)
@@ -394,7 +195,6 @@ def main():
             for t in range(1, args.t_max + 1):
                 gc.collect(); torch.cuda.empty_cache()
 
-                # --- Generate y_t ---
                 humanomni_model.eval()
                 inputs = get_humanomni_inputs(humanomni_processor, generation_video, humanomni_query, sample, args.humanomni_gpu)
 
@@ -405,26 +205,11 @@ def main():
                 y_t_text = humanomni_processor.decode(generated_sequence, skip_special_tokens=True)
                 print(f"  [Iter {t}/{args.t_max}] Answer = {y_t_text}")
 
-                # --- Reviser scores r_t (LoRA disabled) ---
                 base_model.thinker.disable_adapters()
                 score_t, reviser_raw = revise_answer(base_model.thinker, base_processor, video_path, query, y_t_text, args.main_gpu)
-                print(f"  [Iter {t}/{args.t_max}] Score = {score_t}/10")
 
                 all_history.append({"iter": t, "answer": y_t_text, "score": score_t, "reviser_raw": reviser_raw})
 
-                if score_t > best_score:
-                    best_score = score_t
-                    best_answer = y_t_text
-                    best_raw_resp = reviser_raw
-
-                if score_t >= args.tau:
-                    print(f"  -> Score {score_t} >= Tau({args.tau}), breaking loop early.")
-                    early_stop = True
-                    break
-
-                if t == args.t_max:
-                    print(f"  -> Reached T_max, stopping.")
-                    break
 
                 # --- RL Update (REINFORCE) ---
                 humanomni_model.train()
@@ -433,18 +218,6 @@ def main():
                 advantage = float(score_t - b)
                 advantage_tensor = torch.tensor([advantage], device=args.humanomni_gpu, dtype=torch.bfloat16)
 
-                concat_ids = torch.cat([inputs.input_ids, generated_sequence.unsqueeze(0)], dim=1)
-                attention_mask = torch.ones_like(concat_ids)
-                labels = torch.cat([
-                    torch.full_like(inputs.input_ids, -100),
-                    generated_sequence.unsqueeze(0)
-                ], dim=1)
-
-                forward_kwargs = {k: v for k, v in inputs.items() if k not in ["input_ids", "attention_mask", "labels"]}
-                forward_kwargs["input_ids"] = concat_ids
-                forward_kwargs["attention_mask"] = attention_mask
-                forward_kwargs["labels"] = labels
-
                 outputs = humanomni_model(**forward_kwargs)
 
                 nll_loss = outputs.loss
@@ -454,67 +227,7 @@ def main():
                 optimizer.step()
 
                 b = args.alpha * b + (1.0 - args.alpha) * score_t
-                print(f"  -> Updated Baseline b = {b:.2f}, Loss = {final_loss.item():.4f} (NLL: {nll_loss.item():.4f}, Adv: {advantage:.2f})")
-
-            # ====== CLEANUP & PERSIST ======
-            sample_lora_dir = os.path.join(args.lora_save_dir, dataset_id)
-            humanomni_model.save_pretrained(sample_lora_dir, selected_adapters=[adapter_name])
-
-            humanomni_model.delete_adapter(adapter_name)
-            adapter_deleted = True
-
-            save_entry = {
-                "id": dataset_id,
-                "query": query,
-                "video": video_path,
-                "use_grounder": use_grounder,
-                "grounded_span": grounded_span,
-                "collector_output": collector_text,
-                "planner_output": planner_raw,
-                "best_score": best_score,
-                "best_answer": best_answer,
-                "history": all_history,
-                "early_stop": early_stop,
-            }
-            with open(args.output_file, 'a', encoding='utf-8') as f:
-                f.write(json.dumps(save_entry, ensure_ascii=False) + "\n")
-        except Exception as exc:
-            error_type, error_message, traceback_text = format_exception_info(exc)
-            print(f"[Skip] Sample [{dataset_id}] failed with {error_type}: {error_message}")
-            append_skipped_sample_log(
-                dataset_id=dataset_id,
-                query=query,
-                video_path=video_path,
-                error_type=error_type,
-                error_message=error_message,
-                traceback_text=traceback_text,
-                chunk_idx=args.chunk_idx,
-            )
-            print("[Skip] Logged to eval_results/skipped_samples_intentbench.log; sample will be retried in future runs.")
-            continue
-        finally:
-            if optimizer is not None:
-                del optimizer
-            if adapter_name and not adapter_deleted:
-                try:
-                    humanomni_model.delete_adapter(adapter_name)
-                except Exception as cleanup_exc:
-                    print(f"[Cleanup Warning] Failed to delete adapter {adapter_name}: {cleanup_exc}")
-            if trim_path and os.path.exists(trim_path):
-                try:
-                    os.remove(trim_path)
-                except OSError as cleanup_exc:
-                    print(f"[Cleanup Warning] Failed to remove temp file {trim_path}: {cleanup_exc}")
-            gc.collect()
-            torch.cuda.empty_cache()
-
-    grounder_proc.stdin.write(json.dumps({"command": "shutdown"}) + "\n")
-    grounder_proc.stdin.flush()
-    grounder_proc.wait()
-
-if __name__ == "__main__":
-    main()
 
-```
+