| import os |
| import torch |
| import time |
| from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessor, LogitsProcessorList |
| from bridge import LatentBridge |
|
|
| class BridgeDecayProcessor(LogitsProcessor): |
| def __init__(self, bridge): |
| self.bridge = bridge |
|
|
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: |
| self.bridge.step_generation() |
| return scores |
|
|
| def load_project_files(directory, max_chars=20000): |
| """ |
| Load source files from a directory. |
| Limits characters to avoid exceeding the context window or VRAM. |
| """ |
| content = "" |
| valid_extensions = ('.py', '.ts', '.tsx', '.java', '.js') |
| |
| for root, dirs, files in os.walk(directory): |
| |
| if any(ignored in root for ignored in ['node_modules', '.git', '__pycache__', 'dist', 'build']): |
| continue |
| |
| for file in files: |
| if file.endswith(valid_extensions): |
| filepath = os.path.join(root, file) |
| try: |
| with open(filepath, 'r', encoding='utf-8') as f: |
| file_content = f.read() |
| if len(content) + len(file_content) > max_chars: |
| content += f"\n--- {file} (TRUNCATED) ---\n" |
| return content |
| content += f"\n--- {file} ---\n{file_content}\n" |
| except Exception: |
| pass |
| return content |
|
|
| def main(): |
| model_name = "Qwen/Qwen3.5-4B" |
| checkpoint_path = "bridge_weights.pt" |
| |
|
|
| target_dir = r".." |
| |
| print(f"[INFO] Loading Base Model: {model_name} in FP16...") |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_name, |
| torch_dtype=torch.float16, |
| attn_implementation="sdpa" |
| ).to("cuda") |
| |
| print(f"[INFO] Loading Latent Bridge...") |
| bridge = LatentBridge(hidden_dim=model.config.hidden_size, target_layers=[11, 19, 27]) |
| ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=True) |
| bridge.load_state_dict(ckpt["state_dict"] if "state_dict" in ckpt else ckpt) |
| device = next(model.parameters()).device |
| bridge = bridge.to(device) |
|
|
|
|
| print(f"\n[INFO] Loading source files from: {os.path.abspath(target_dir)}...") |
| project_code = load_project_files(target_dir, max_chars=25000) |
| print(f"[INFO] Code loaded: ~{len(project_code)} characters.") |
|
|
| sys_A = "System: You are a Senior Software Architect. Explain the architecture of the provided code clearly and schematically. Describe how the various components interact. Do not think out loud, just produce the final report." |
| sys_B = "System: You are an expert code reviewer (Staff Engineer). Deeply study the code, analyze data flows, class responsibilities, and memorize business logic and architectural interactions. Build a complex mental map." |
| |
| user_prompt = f"User: Analyze the following source code and explain its architecture in detail.\n\nCodebase:\n{project_code}\n\nWhat is the architecture of the system?" |
|
|
|
|
| msgs_A = [{"role": "system", "content": sys_A}, {"role": "user", "content": user_prompt}] |
| msgs_B = [{"role": "system", "content": sys_B}, {"role": "user", "content": user_prompt}] |
| |
| prompt_A = tokenizer.apply_chat_template(msgs_A, tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| prompt_B = tokenizer.apply_chat_template(msgs_B, tokenize=False, add_generation_prompt=True, enable_thinking=False) |
|
|
| inputs_A = tokenizer(prompt_A, return_tensors="pt").to(device) |
| inputs_B = tokenizer(prompt_B, return_tensors="pt").to(device) |
|
|
| print("\n[INFO] === PHASE 1: Background Architecture Study (Agent B) ===") |
| bridge.detach() |
| bridge.clear_context() |
| with torch.no_grad(): |
| outputs_B = model(**inputs_B, output_hidden_states=True) |
| bridge.set_context(outputs_B.hidden_states) |
| print("[SUCCESS] Mental map of the code captured in Latent Space!") |
|
|
| print("\n[INFO] === PHASE 2: Explanation (Agent A guided by Latent Bridge) ===") |
| bridge.enable_generation_mode(decay_rate=0.85) |
| bridge.attach(model) |
| |
| processors = LogitsProcessorList([BridgeDecayProcessor(bridge)]) |
|
|
| start_time = time.time() |
| with torch.no_grad(): |
| out = model.generate( |
| **inputs_A, |
| max_new_tokens=4096, |
| do_sample=False, |
| repetition_penalty=1.1, |
| logits_processor=processors, |
| pad_token_id=tokenizer.eos_token_id |
| ) |
| end_time = time.time() |
| |
| bridge.disable_generation_mode() |
| bridge.detach() |
|
|
| new_tokens = out[0][inputs_A.input_ids.shape[1]:] |
| generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True) |
| |
| generated_text = generated_text.strip() |
|
|
| print(f"\n[SUCCESS] Generated in {end_time - start_time:.2f} sec.") |
| print("\nARCHITECTURAL ANALYSIS (4B with Telepathy):\n" + "="*80) |
| print(generated_text) |
| print("="*80) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
