File size: 5,128 Bytes
3070f60 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | 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):
# Ignore useless directories
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 # Limit reached
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()
|