Delete modeling_tx_standalone.py

modeling_tx_standalone.py

@@ -1,157 +0,0 @@
-# Copyright (C) Tahoe Therapeutics 2025. All rights reserved.
-"""
-HuggingFace-compatible wrapper for TXModel (Standalone version)
-Only requires: transformers, torch, safetensors
-"""
-
-from typing import Optional, Union, Tuple
-import torch
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutput
-
-from configuration_tx import TXConfig
-from model_standalone import TXModel
-
-
-class TXPreTrainedModel(PreTrainedModel):
-    """
-    Base class for TXModel with HuggingFace integration
-    """
-    config_class = TXConfig
-    base_model_prefix = "tx_model"
-    supports_gradient_checkpointing = False
-    _no_split_modules = ["TXBlock"]
-    
-    def _init_weights(self, module):
-        """Initialize weights"""
-        if isinstance(module, torch.nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=0.02)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, torch.nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=0.02)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, torch.nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-
-class TXModelForHF(TXPreTrainedModel):
-    """
-    HuggingFace-compatible TXModel
-    
-    This model can be used directly with HuggingFace's from_pretrained()
-    and requires only: transformers, torch, safetensors
-    
-    No dependencies on llmfoundry, composer, or other external libraries.
-    """
-    
-    def __init__(self, config: TXConfig):
-        super().__init__(config)
-        
-        # Initialize standalone model
-        self.tx_model = TXModel(
-            vocab_size=config.vocab_size,
-            d_model=config.d_model,
-            n_layers=config.n_layers,
-            n_heads=config.n_heads,
-            expansion_ratio=config.expansion_ratio,
-            pad_token_id=config.pad_token_id,
-            pad_value=config.pad_value,
-            num_bins=config.num_bins,
-            norm_scheme=config.norm_scheme,
-            transformer_activation=config.transformer_activation,
-            cell_emb_style=config.cell_emb_style,
-            use_chem_token=config.use_chem_token,
-            attn_config=config.attn_config,
-            norm_config=config.norm_config,
-            gene_encoder_config=config.gene_encoder_config,
-            expression_encoder_config=config.expression_encoder_config,
-            expression_decoder_config=config.expression_decoder_config,
-            mvc_config=config.mvc_config,
-            chemical_encoder_config=config.chemical_encoder_config,
-            use_glu=config.use_glu,
-            return_gene_embeddings=config.return_gene_embeddings,
-            keep_first_n_tokens=config.keep_first_n_tokens,
-        )
-        
-        # Post init
-        self.post_init()
-    
-    def forward(
-        self,
-        genes: torch.Tensor,
-        values: torch.Tensor,
-        gen_masks: torch.Tensor,
-        key_padding_mask: Optional[torch.Tensor] = None,
-        drug_ids: Optional[torch.Tensor] = None,
-        skip_decoders: bool = False,
-        output_hidden_states: bool = False,
-        return_dict: bool = True,
-    ) -> Union[Tuple, BaseModelOutput]:
-        """
-        Forward pass through the model.
-        
-        Args:
-            genes: Gene token IDs [batch_size, seq_len]
-            values: Expression values [batch_size, seq_len]
-            gen_masks: Generation masks [batch_size, seq_len]
-            key_padding_mask: Padding mask [batch_size, seq_len]
-            drug_ids: Drug IDs [batch_size] (optional)
-            skip_decoders: Whether to skip decoder computation
-            output_hidden_states: Whether to return hidden states
-            return_dict: Whether to return a dict or tuple
-            
-        Returns:
-            Model outputs
-        """
-        
-        if key_padding_mask is None:
-            key_padding_mask = ~genes.eq(self.config.pad_token_id)
-        
-        outputs = self.tx_model(
-            genes=genes,
-            values=values,
-            gen_masks=gen_masks,
-            key_padding_mask=key_padding_mask,
-            drug_ids=drug_ids,
-            skip_decoders=skip_decoders,
-            output_hidden_states=output_hidden_states,
-        )
-        
-        if not return_dict:
-            return tuple(v for v in outputs.values())
-        
-        # Convert to HuggingFace output format
-        return BaseModelOutput(
-            last_hidden_state=outputs.get("cell_emb"),
-            hidden_states=outputs.get("hidden_states") if output_hidden_states else None,
-        )
-    
-    def get_input_embeddings(self):
-        """Get input embeddings"""
-        return self.tx_model.gene_encoder.embedding
-    
-    def set_input_embeddings(self, value):
-        """Set input embeddings"""
-        self.tx_model.gene_encoder.embedding = value
-    
-    def get_output_embeddings(self):
-        """Get output embeddings (not applicable)"""
-        return None
-    
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        """
-        Load model from pretrained weights.
-        
-        Works with both local paths and HuggingFace Hub.
-        Requires only: transformers, torch, safetensors
-        """
-        # Let parent class handle config and weight loading
-        return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
-
-
-# Alias for easier importing
-TXForCausalLM = TXModelForHF