Create alm_qwen.py

alm_qwen.py

@@ -0,0 +1,347 @@
+# --- START OF FILE alm_qwen_hf.py ---
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from typing import List, Tuple, Dict, Any
+import os
+import json
+
+# Assuming ALM.py is in the same directory or accessible in PYTHONPATH
+from ALM import AttentionLinkedMemory # Make sure ALM.py is saved and accessible
+
+class QwenGenerator(nn.Module):
+    def __init__(self, model_name_or_path: str, device="cuda", tokenizer_path: str = None):
+        super().__init__()
+        self.device = device
+        self.model_name_or_path = model_name_or_path # Store for saving config
+        self.tokenizer_path = tokenizer_path if tokenizer_path else model_name_or_path
+
+        print(f"Loading Qwen model from: {self.model_name_or_path}...")
+        print(f"Loading Qwen tokenizer from: {self.tokenizer_path}...")
+
+        # Standard loading (requires more resources)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name_or_path,
+            torch_dtype="auto",
+            device_map="auto",
+            trust_remote_code=True
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path, trust_remote_code=True)
+        
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "left"
+
+        print(f"Qwen model and tokenizer loaded. Model device: {self.model.device}")
+
+    def format_prompt(self, query: str, context_snippets: List[str]) -> str:
+        if context_snippets:
+            context_str = "\n".join(f"- {cs}" for cs in context_snippets)
+            # Qwen specific chat format
+            final_prompt_str = "<|im_start|>system\nYou are a helpful assistant. Use the provided context to answer the user's query. If the context is insufficient, say so.<|im_end|>\n"
+            final_prompt_str += "<|im_start|>user\n"
+            final_prompt_str += f"Context:\n{context_str}\n\n"
+            final_prompt_str += f"Query:\n{query}\n<|im_end|>\n<|im_start|>assistant\n"
+        else:
+            # Qwen specific chat format without context
+            final_prompt_str = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+            final_prompt_str += "<|im_start|>user\n"
+            final_prompt_str += f"Query:\n{query}\n<|im_end|>\n<|im_start|>assistant\n"
+        return final_prompt_str
+
+    def generate(self, prompts: List[str], max_new_tokens: int = 150, **kwargs) -> List[str]:
+        self.model.eval()
+        inputs = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, max_length=2048)
+        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                do_sample=kwargs.get("do_sample", True),
+                temperature=kwargs.get("temperature", 0.7),
+                top_p=kwargs.get("top_p", 0.9),
+                **kwargs
+            )
+        
+        decoded_outputs = []
+        for i, output_ids in enumerate(outputs):
+            prompt_len = inputs['input_ids'][i].shape[0]
+            generated_ids = output_ids[prompt_len:]
+            decoded_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
+            decoded_outputs.append(decoded_text.strip())
+            
+        return decoded_outputs
+
+    def save_pretrained(self, save_directory: str):
+        """Saves the Qwen model and tokenizer to a directory."""
+        model_save_path = os.path.join(save_directory, "qwen_model")
+        tokenizer_save_path = os.path.join(save_directory, "qwen_tokenizer")
+        
+        print(f"Saving Qwen model to {model_save_path}")
+        self.model.save_pretrained(model_save_path)
+        print(f"Saving Qwen tokenizer to {tokenizer_save_path}")
+        self.tokenizer.save_pretrained(tokenizer_save_path)
+
+class ALMQwenModel_HF(nn.Module):
+    def __init__(self,
+                 alm_config: Dict[str, Any],
+                 qwen_model_name_or_path: str, # Can be HF name or local path
+                 qwen_tokenizer_path: str = None, # Optional separate path for tokenizer
+                 device: str = "cuda" if torch.cuda.is_available() else "cpu",
+                 top_k_buckets: int = 3,
+                 top_k_items_per_bucket: int = 2):
+        super().__init__()
+        self.device = device
+        self.alm_config = alm_config # Store for saving
+        self.qwen_model_name_or_path = qwen_model_name_or_path # Store for saving
+        self.qwen_tokenizer_path = qwen_tokenizer_path
+        self.top_k_buckets = top_k_buckets # Store for saving
+        self.top_k_items_per_bucket = top_k_items_per_bucket # Store for saving
+
+        self.alm_layer = AttentionLinkedMemory(**alm_config).to(device)
+        self.qwen_generator = QwenGenerator(
+            model_name_or_path=qwen_model_name_or_path,
+            device=device,
+            tokenizer_path=qwen_tokenizer_path
+        )
+
+    def forward(self,
+                query_texts: List[str],
+                query_embeddings_for_alm: torch.Tensor,
+                memory_item_embeddings: torch.Tensor,
+                memory_text_items: List[List[List[str]]],
+                memory_mask: torch.Tensor = None
+                ) -> Tuple[List[str], torch.Tensor, torch.Tensor]:
+        self.alm_layer.eval()
+        query_embeddings_for_alm = query_embeddings_for_alm.to(self.device)
+        memory_item_embeddings = memory_item_embeddings.to(self.device)
+        if memory_mask is not None:
+            memory_mask = memory_mask.to(self.device)
+
+        with torch.no_grad():
+            _, bucket_att_weights, item_att_weights = self.alm_layer(
+                query_embeddings_for_alm, memory_item_embeddings, memory_mask
+            )
+        
+        batch_retrieved_texts: List[List[str]] = []
+        for b_idx in range(len(query_texts)):
+            retrieved_for_sample: List[str] = []
+            current_bucket_weights = bucket_att_weights[b_idx]
+            _, top_bucket_indices = torch.topk(current_bucket_weights,
+                                               k=min(self.top_k_buckets, current_bucket_weights.size(0)))
+
+            for bucket_idx in top_bucket_indices:
+                bucket_idx_item = bucket_idx.item()
+                current_item_weights = item_att_weights[b_idx, bucket_idx_item, :]
+                
+                if memory_mask is not None:
+                    item_m = memory_mask[b_idx, bucket_idx_item, :]
+                    current_item_weights = current_item_weights.masked_fill(item_m == 0, -float('inf'))
+                
+                num_valid_items = (current_item_weights > -float('inf')).sum().item()
+                if num_valid_items == 0: continue
+
+                _, top_item_indices_in_bucket = torch.topk(current_item_weights,
+                                                           k=min(self.top_k_items_per_bucket, num_valid_items))
+                
+                for item_idx_in_bucket in top_item_indices_in_bucket:
+                    item_idx_in_bucket_item = item_idx_in_bucket.item()
+                    if memory_mask is not None and not memory_mask[b_idx, bucket_idx_item, item_idx_in_bucket_item]:
+                        continue
+                    try:
+                        text_content = memory_text_items[b_idx][bucket_idx_item][item_idx_in_bucket_item]
+                        if text_content:
+                             retrieved_for_sample.append(text_content)
+                    except IndexError:
+                        print(f"Warning: IndexError accessing memory_text_items[{b_idx}][{bucket_idx_item}][{item_idx_in_bucket_item}]")
+                        continue
+            batch_retrieved_texts.append(list(dict.fromkeys(retrieved_for_sample)))
+
+        prompts_for_qwen = []
+        for i, q_text in enumerate(query_texts):
+            prompt = self.qwen_generator.format_prompt(q_text, batch_retrieved_texts[i])
+            prompts_for_qwen.append(prompt)
+        
+        generated_answers = self.qwen_generator.generate(prompts_for_qwen)
+        return generated_answers, bucket_att_weights, item_att_weights
+
+    def save_model(self, save_directory: str):
+        """Saves the entire ALMQwenModel_HF to the specified directory."""
+        os.makedirs(save_directory, exist_ok=True)
+
+        # 1. Save ALM layer state_dict
+        alm_state_dict_path = os.path.join(save_directory, "alm_layer_state_dict.pth")
+        torch.save(self.alm_layer.state_dict(), alm_state_dict_path)
+        print(f"ALM layer state_dict saved to {alm_state_dict_path}")
+
+        # 2. Save QwenGenerator (model and tokenizer)
+        qwen_save_path = os.path.join(save_directory, "qwen_generator")
+        os.makedirs(qwen_save_path, exist_ok=True)
+        self.qwen_generator.save_pretrained(qwen_save_path)
+        print(f"Qwen generator (model & tokenizer) saved in {qwen_save_path}")
+
+        # 3. Save ALMQwenModel_HF configurations
+        config = {
+            "alm_config": self.alm_config,
+            # Store relative paths for qwen model/tokenizer for portability
+            "qwen_model_name_or_path": "qwen_generator/qwen_model", # Relative path
+            "qwen_tokenizer_path": "qwen_generator/qwen_tokenizer", # Relative path
+            "top_k_buckets": self.top_k_buckets,
+            "top_k_items_per_bucket": self.top_k_items_per_bucket
+        }
+        config_path = os.path.join(save_directory, "alm_qwen_hf_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(config, f, indent=4)
+        print(f"ALMQwenModel_HF configuration saved to {config_path}")
+        print(f"Model saved successfully to {save_directory}")
+
+    @classmethod
+    def load_model(cls, load_directory: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
+        """Loads an ALMQwenModel_HF from the specified directory."""
+        print(f"Loading model from {load_directory}...")
+
+        # 1. Load ALMQwenModel_HF configurations
+        config_path = os.path.join(load_directory, "alm_qwen_hf_config.json")
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"Configuration file not found: {config_path}")
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+        
+        alm_config = config["alm_config"]
+        # Construct absolute paths for qwen model and tokenizer from saved relative paths
+        qwen_model_path = os.path.join(load_directory, config["qwen_model_name_or_path"])
+        qwen_tokenizer_path = os.path.join(load_directory, config["qwen_tokenizer_path"])
+        top_k_buckets = config["top_k_buckets"]
+        top_k_items_per_bucket = config["top_k_items_per_bucket"]
+
+        # 2. Instantiate the model (QwenGenerator will load its components from the paths)
+        model = cls(
+            alm_config=alm_config,
+            qwen_model_name_or_path=qwen_model_path,
+            qwen_tokenizer_path=qwen_tokenizer_path,
+            device=device,
+            top_k_buckets=top_k_buckets,
+            top_k_items_per_bucket=top_k_items_per_bucket
+        )
+        print("ALMQwenModel_HF structure initialized.")
+
+        # 3. Load ALM layer state_dict
+        alm_state_dict_path = os.path.join(load_directory, "alm_layer_state_dict.pth")
+        if not os.path.exists(alm_state_dict_path):
+            raise FileNotFoundError(f"ALM state_dict not found: {alm_state_dict_path}")
+        
+        # Ensure the model's ALM layer is on the correct device before loading state_dict
+        model.alm_layer.to(device)
+        state_dict = torch.load(alm_state_dict_path, map_location=device)
+        model.alm_layer.load_state_dict(state_dict)
+        print(f"ALM layer state_dict loaded from {alm_state_dict_path}")
+        
+        # Qwen model is already loaded by QwenGenerator instantiation on the correct device due to device_map="auto"
+        # or manually if we passed device to QwenGenerator more directly.
+        # If device_map="auto" was used, it might be on multiple devices.
+        # Ensure the overall model object has its device attribute set.
+        model.device = device # Ensure the main model object knows its primary device.
+        
+        print(f"Model loaded successfully from {load_directory} and placed on device: {device}")
+        return model
+
+
+# ========================= Example Usage for ALM-Qwen with Hugging Face =========================
+if __name__ == "__main__":
+    print("\n--- Testing ALM-Qwen with Hugging Face Qwen ---")
+    
+    _device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {_device}")
+
+    # --- Hyperparameters ---
+    _batch_size = 1 # Reduced for faster testing
+    _alm_query_dim = 128
+    _alm_memory_dim = 64
+    _alm_embed_dim = 256
+    _alm_num_heads = 8
+    _alm_output_dim = 128
+    _num_kb_buckets = 3
+    _max_kb_items_per_bucket = 5
+
+    _alm_config_example = {
+        'query_dim': _alm_query_dim,
+        'memory_dim': _alm_memory_dim,
+        'embed_dim': _alm_embed_dim,
+        'num_heads': _alm_num_heads,
+        'output_dim': _alm_output_dim,
+        'dropout_rate': 0.0
+    }
+
+    _query_texts_for_qwen = ["What is attention in LLMs?"]
+    _query_embeddings_for_alm = torch.randn(_batch_size, _alm_query_dim)
+    _kb_memory_item_embeddings = torch.randn(_batch_size, _num_kb_buckets, _max_kb_items_per_bucket, _alm_memory_dim)
+    _kb_memory_text_items: List[List[List[str]]] = []
+    for b in range(_batch_size):
+        batch_sample_text = []
+        for i in range(_num_kb_buckets):
+            bucket_texts = [f"Doc {b+1}-B{i+1}-I{j+1}: info snippet {j}." for j in range(_max_kb_items_per_bucket)]
+            batch_sample_text.append(bucket_texts)
+        _kb_memory_text_items.append(batch_sample_text)
+    _kb_memory_mask = torch.ones(_batch_size, _num_kb_buckets, _max_kb_items_per_bucket, dtype=torch.bool)
+    _kb_memory_mask[:, :, -1:] = False # Mask last item
+
+    _qwen_model_name = "Qwen/Qwen2.5-0.5B-Instruct" # Smaller model for testing
+
+    try:
+        # --- Instantiate Original Model ---
+        print("\n--- Creating and testing original model ---")
+        original_model = ALMQwenModel_HF(
+            alm_config=_alm_config_example,
+            qwen_model_name_or_path=_qwen_model_name,
+            device=_device,
+            top_k_buckets=2,
+            top_k_items_per_bucket=1
+        )
+        
+        # Optional: Dummy forward pass to ensure everything is initialized
+        # (especially lazy initializations if any, though not typical here)
+        _ = original_model(
+             _query_texts_for_qwen, _query_embeddings_for_alm, _kb_memory_item_embeddings,
+             _kb_memory_text_items, _kb_memory_mask
+        )
+        print("Original model created and tested with a dummy pass.")
+
+        # --- Save the Model ---
+        save_dir = "./saved_alm_qwen_model"
+        print(f"\n--- Saving model to {save_dir} ---")
+        original_model.save_model(save_dir)
+
+        # --- Load the Model ---
+        print(f"\n--- Loading model from {save_dir} ---")
+        # Ensure to pass the target device for loading
+        loaded_model = ALMQwenModel_HF.load_model(save_dir, device=_device)
+        print("Model loaded successfully.")
+
+        # --- Test Loaded Model ---
+        print("\n--- Testing loaded model ---")
+        generated_answers, _, _ = loaded_model(
+            _query_texts_for_qwen,
+            _query_embeddings_for_alm,
+            _kb_memory_item_embeddings,
+            _kb_memory_text_items,
+            _kb_memory_mask
+        )
+        print("\n--- Results from Loaded Model ---")
+        for i in range(len(_query_texts_for_qwen)):
+            print(f"Query {i+1}: {_query_texts_for_qwen[i]}")
+            print(f"  Generated Answer {i+1}: {generated_answers[i]}")
+            print("-" * 30)
+        
+        print("\nSave and Load test completed.")
+
+    except ImportError as e:
+        print(f"ImportError: {e}.")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        import traceback
+        traceback.print_exc()
+
+# --- END OF FILE alm_qwen_hf.py ---