Upload folder using huggingface_hub

README.md

@@ -0,0 +1,88 @@
+---
+language:
+- en
+tags:
+- biology
+- dna
+- genomics
+- metagenomics
+- classifier
+- awd-lstm
+- transfer-learning
+license: mit
+pipeline_tag: text-classification
+library_name: pytorch
+---
+
+# LookingGlass Functional Classifier
+
+Classifies DNA reads into one of 1274 experimentally-validated functional annotations with 81.5% accuracy.
+
+This is a **pure PyTorch implementation** fine-tuned from the LookingGlass base model.
+
+## Links
+
+- **Paper**: [Deep learning of a bacterial and archaeal universal language of life enables transfer learning and illuminates microbial dark matter](https://doi.org/10.1038/s41467-022-30070-8) (Nature Communications, 2022)
+- **GitHub**: [ahoarfrost/LookingGlass](https://github.com/ahoarfrost/LookingGlass)
+- **Base Model**: [HoarfrostLab/lookingglass-v1](https://huggingface.co/HoarfrostLab/lookingglass-v1)
+
+## Citation
+
+```bibtex
+@article{hoarfrost2022deep,
+  title={Deep learning of a bacterial and archaeal universal language of life
+         enables transfer learning and illuminates microbial dark matter},
+  author={Hoarfrost, Adrienne and Aptekmann, Ariel and Farfanuk, Gaetan and Bromberg, Yana},
+  journal={Nature Communications},
+  volume={13},
+  number={1},
+  pages={2606},
+  year={2022},
+  publisher={Nature Publishing Group}
+}
+```
+
+## Model
+
+| | |
+|---|---|
+| Architecture | LookingGlass encoder + classification head |
+| Encoder | AWD-LSTM (3-layer, unidirectional) |
+| Classes | 1274 functional annotation classes |
+| Parameters | ~17M |
+
+## Installation
+
+```bash
+pip install torch
+git clone https://huggingface.co/HoarfrostLab/LGv1_FunctionalClassifier
+cd LGv1_FunctionalClassifier
+```
+
+## Usage
+
+```python
+from lookingglass_classifier import LookingGlassClassifier, LookingGlassTokenizer
+
+model = LookingGlassClassifier.from_pretrained('.')
+tokenizer = LookingGlassTokenizer()
+model.eval()
+
+inputs = tokenizer(["GATTACA", "ATCGATCGATCG"], return_tensors=True)
+
+# Get predictions
+predictions = model.predict(inputs['input_ids'])
+print(predictions)  # tensor([class_idx, class_idx])
+
+# Get probabilities
+probs = model.predict_proba(inputs['input_ids'])
+print(probs.shape)  # torch.Size([2, 1274])
+
+# Get raw logits
+logits = model(inputs['input_ids'])
+print(logits.shape)  # torch.Size([2, 1274])
+```
+
+## License
+
+MIT License

config.json

@@ -0,0 +1,22 @@
+{
+  "vocab_size": 8,
+  "hidden_size": 104,
+  "intermediate_size": 1152,
+  "num_hidden_layers": 3,
+  "pad_token_id": 1,
+  "bos_token_id": 2,
+  "eos_token_id": 3,
+  "bidirectional": false,
+  "output_dropout": 0.1,
+  "hidden_dropout": 0.15,
+  "input_dropout": 0.25,
+  "embed_dropout": 0.02,
+  "weight_dropout": 0.2,
+  "tie_weights": true,
+  "output_bias": true,
+  "model_type": "lookingglass",
+  "num_classes": 1274,
+  "classifier_hidden": 50,
+  "classifier_dropout": 0.0,
+  "class_names": []
+}

lookingglass.py

@@ -0,0 +1,780 @@
+"""
+LookingGlass - A DNA Language Model
+
+Pure PyTorch implementation of LookingGlass, a pretrained language model for DNA sequences.
+Based on AWD-LSTM architecture, originally trained with fastai v1.
+
+Paper: Hoarfrost et al., "Deep learning of a bacterial and archaeal universal language
+of life enables transfer learning and illuminates microbial dark matter",
+Nature Communications, 2022.
+
+Usage:
+    from lookingglass import LookingGlass, LookingGlassTokenizer
+
+    # Load from HuggingFace Hub
+    model = LookingGlass.from_pretrained('HoarfrostLab/lookingglass-v1')
+    tokenizer = LookingGlassTokenizer()
+
+    # Or load from local path
+    model = LookingGlass.from_pretrained('./lookingglass-v1')
+
+    inputs = tokenizer(["GATTACA", "ATCGATCG"], return_tensors=True)
+    embeddings = model.get_embeddings(inputs['input_ids'])  # (batch, 104)
+"""
+
+import json
+import os
+import warnings
+from dataclasses import dataclass, asdict
+from typing import Optional, Tuple, List, Dict, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from huggingface_hub import hf_hub_download
+    HF_HUB_AVAILABLE = True
+except ImportError:
+    HF_HUB_AVAILABLE = False
+
+
+__version__ = "1.1.0"
+
+
+def _is_hf_hub_id(path: str) -> bool:
+    """Check if path looks like a HuggingFace Hub model ID (e.g., 'user/model')."""
+    if os.path.exists(path):
+        return False
+    return '/' in path and not path.startswith(('.', '/'))
+
+
+def _download_from_hub(repo_id: str, filename: str) -> str:
+    """Download a file from HuggingFace Hub and return the local path."""
+    if not HF_HUB_AVAILABLE:
+        raise ImportError(
+            "huggingface_hub is required to load models from the Hub. "
+            "Install it with: pip install huggingface_hub"
+        )
+    return hf_hub_download(repo_id=repo_id, filename=filename)
+__all__ = [
+    "LookingGlassConfig",
+    "LookingGlass",
+    "LookingGlassLM",
+    "LookingGlassTokenizer",
+]
+
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+@dataclass
+class LookingGlassConfig:
+    """
+    Configuration for LookingGlass model.
+
+    Default values match the original pretrained LookingGlass model.
+    """
+    vocab_size: int = 8
+    hidden_size: int = 104          # embedding/output size
+    intermediate_size: int = 1152   # LSTM hidden size
+    num_hidden_layers: int = 3
+    pad_token_id: int = 1
+    bos_token_id: int = 2
+    eos_token_id: int = 3
+    bidirectional: bool = False     # original LG is unidirectional
+    output_dropout: float = 0.1
+    hidden_dropout: float = 0.15
+    input_dropout: float = 0.25
+    embed_dropout: float = 0.02
+    weight_dropout: float = 0.2
+    tie_weights: bool = True
+    output_bias: bool = True
+    model_type: str = "lookingglass"
+
+    def to_dict(self) -> Dict:
+        return asdict(self)
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+        with open(os.path.join(save_directory, "config.json"), 'w') as f:
+            json.dump(self.to_dict(), f, indent=2)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_path: str) -> "LookingGlassConfig":
+        if _is_hf_hub_id(pretrained_path):
+            try:
+                config_path = _download_from_hub(pretrained_path, "config.json")
+            except Exception:
+                return cls()
+        elif os.path.isdir(pretrained_path):
+            config_path = os.path.join(pretrained_path, "config.json")
+        else:
+            config_path = pretrained_path
+
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                config_dict = json.load(f)
+            valid_fields = {f.name for f in cls.__dataclass_fields__.values()}
+            return cls(**{k: v for k, v in config_dict.items() if k in valid_fields})
+        return cls()
+
+
+# =============================================================================
+# Tokenizer
+# =============================================================================
+
+VOCAB = ['xxunk', 'xxpad', 'xxbos', 'xxeos', 'G', 'A', 'C', 'T']
+VOCAB_TO_ID = {tok: i for i, tok in enumerate(VOCAB)}
+ID_TO_VOCAB = {i: tok for i, tok in enumerate(VOCAB)}
+
+
+class LookingGlassTokenizer:
+    """
+    Tokenizer for DNA sequences.
+
+    Each nucleotide (G, A, C, T) is a single token. By default, adds BOS token
+    at the start of each sequence (matching original LookingGlass training).
+
+    Special tokens:
+        - xxunk (0): Unknown
+        - xxpad (1): Padding
+        - xxbos (2): Beginning of sequence
+        - xxeos (3): End of sequence
+    """
+
+    vocab = VOCAB
+    vocab_to_id = VOCAB_TO_ID
+    id_to_vocab = ID_TO_VOCAB
+
+    def __init__(
+        self,
+        add_bos_token: bool = True,   # original LG uses BOS
+        add_eos_token: bool = False,  # original LG does not use EOS
+        padding_side: str = "right",
+    ):
+        self.add_bos_token = add_bos_token
+        self.add_eos_token = add_eos_token
+        self.padding_side = padding_side
+
+        self.unk_token_id = 0
+        self.pad_token_id = 1
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    def encode(self, sequence: str, add_special_tokens: bool = True) -> List[int]:
+        """Encode a DNA sequence to token IDs."""
+        tokens = []
+
+        if add_special_tokens and self.add_bos_token:
+            tokens.append(self.bos_token_id)
+
+        for char in sequence.upper():
+            if char in self.vocab_to_id:
+                tokens.append(self.vocab_to_id[char])
+            elif char.strip():
+                tokens.append(self.unk_token_id)
+
+        if add_special_tokens and self.add_eos_token:
+            tokens.append(self.eos_token_id)
+
+        return tokens
+
+    def decode(self, token_ids: Union[List[int], torch.Tensor], skip_special_tokens: bool = True) -> str:
+        """Decode token IDs back to DNA sequence."""
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+
+        special_ids = {0, 1, 2, 3}
+        tokens = []
+        for tid in token_ids:
+            if skip_special_tokens and tid in special_ids:
+                continue
+            tokens.append(self.id_to_vocab.get(tid, 'xxunk'))
+        return ''.join(tokens)
+
+    def __call__(
+        self,
+        sequences: Union[str, List[str]],
+        padding: Union[bool, str] = False,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        return_tensors: Union[bool, str] = False,
+        return_attention_mask: bool = True,
+    ) -> Dict[str, torch.Tensor]:
+        """Tokenize DNA sequence(s)."""
+        if isinstance(sequences, str):
+            sequences = [sequences]
+            single = True
+        else:
+            single = False
+
+        encoded = [self.encode(seq) for seq in sequences]
+
+        if truncation and max_length:
+            encoded = [e[:max_length] for e in encoded]
+
+        # Padding
+        if padding or len(encoded) > 1:
+            if padding == 'max_length' and max_length:
+                pad_len = max_length
+            else:
+                pad_len = max(len(e) for e in encoded)
+
+            padded = []
+            masks = []
+            for e in encoded:
+                pad_amount = pad_len - len(e)
+                mask = [1] * len(e) + [0] * pad_amount
+                if self.padding_side == 'right':
+                    e = e + [self.pad_token_id] * pad_amount
+                else:
+                    e = [self.pad_token_id] * pad_amount + e
+                    mask = [0] * pad_amount + [1] * len(e)
+                padded.append(e)
+                masks.append(mask)
+            encoded = padded
+        else:
+            masks = [[1] * len(e) for e in encoded]
+
+        result = {}
+        if return_tensors in ('pt', True):
+            result['input_ids'] = torch.tensor(encoded, dtype=torch.long)
+            if return_attention_mask:
+                result['attention_mask'] = torch.tensor(masks, dtype=torch.long)
+        else:
+            result['input_ids'] = encoded[0] if single else encoded
+            if return_attention_mask:
+                result['attention_mask'] = masks[0] if single else masks
+
+        return result
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+        with open(os.path.join(save_directory, "vocab.json"), 'w') as f:
+            json.dump(self.vocab_to_id, f, indent=2)
+        with open(os.path.join(save_directory, "tokenizer_config.json"), 'w') as f:
+            json.dump({
+                "add_bos_token": self.add_bos_token,
+                "add_eos_token": self.add_eos_token,
+                "padding_side": self.padding_side,
+            }, f, indent=2)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_path: str) -> "LookingGlassTokenizer":
+        kwargs = {}
+        if _is_hf_hub_id(pretrained_path):
+            try:
+                config_path = _download_from_hub(pretrained_path, "tokenizer_config.json")
+                with open(config_path, 'r') as f:
+                    kwargs = json.load(f)
+            except Exception:
+                pass
+        else:
+            config_path = os.path.join(pretrained_path, "tokenizer_config.json")
+            if os.path.exists(config_path):
+                with open(config_path, 'r') as f:
+                    kwargs = json.load(f)
+        return cls(**kwargs)
+
+
+# =============================================================================
+# Model Components
+# =============================================================================
+
+def _dropout_mask(x: torch.Tensor, size: Tuple[int, ...], p: float) -> torch.Tensor:
+    """Create dropout mask with inverted scaling."""
+    return x.new_empty(*size).bernoulli_(1 - p).div_(1 - p)
+
+
+class _RNNDropout(nn.Module):
+    """Dropout consistent across sequence dimension."""
+
+    def __init__(self, p: float = 0.5):
+        super().__init__()
+        self.p = p
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.training or self.p == 0.:
+            return x
+        mask = _dropout_mask(x.data, (x.size(0), 1, x.size(2)), self.p)
+        return x * mask
+
+
+class _EmbeddingDropout(nn.Module):
+    """Dropout applied to entire embedding rows."""
+
+    def __init__(self, embedding: nn.Embedding, p: float):
+        super().__init__()
+        self.embedding = embedding
+        self.p = p
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training and self.p != 0:
+            mask = _dropout_mask(self.embedding.weight.data,
+                                 (self.embedding.weight.size(0), 1), self.p)
+            masked_weight = self.embedding.weight * mask
+        else:
+            masked_weight = self.embedding.weight
+
+        padding_idx = self.embedding.padding_idx if self.embedding.padding_idx is not None else -1
+        return F.embedding(x, masked_weight, padding_idx,
+                          self.embedding.max_norm, self.embedding.norm_type,
+                          self.embedding.scale_grad_by_freq, self.embedding.sparse)
+
+
+class _WeightDropout(nn.Module):
+    """DropConnect applied to RNN hidden-to-hidden weights."""
+
+    def __init__(self, module: nn.Module, p: float, layer_names='weight_hh_l0'):
+        super().__init__()
+        self.module = module
+        self.p = p
+        self.layer_names = [layer_names] if isinstance(layer_names, str) else layer_names
+
+        for layer in self.layer_names:
+            w = getattr(self.module, layer)
+            delattr(self.module, layer)
+            self.register_parameter(f'{layer}_raw', nn.Parameter(w.data))
+            setattr(self.module, layer, w.clone())
+
+        if isinstance(self.module, nn.RNNBase):
+            self.module.flatten_parameters = lambda: None
+
+    def _set_weights(self):
+        for layer in self.layer_names:
+            raw_w = getattr(self, f'{layer}_raw')
+            w = F.dropout(raw_w, p=self.p, training=self.training) if self.training else raw_w.clone()
+            setattr(self.module, layer, w)
+
+    def forward(self, *args):
+        self._set_weights()
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=UserWarning)
+            return self.module(*args)
+
+
+class _AWDLSTMEncoder(nn.Module):
+    """AWD-LSTM encoder backbone."""
+
+    _init_range = 0.1
+
+    def __init__(self, config: LookingGlassConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.num_layers = config.num_hidden_layers
+        self.num_directions = 2 if config.bidirectional else 1
+        self._batch_size = 1
+
+        # Embedding
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size,
+                                         padding_idx=config.pad_token_id)
+        self.embed_tokens.weight.data.uniform_(-self._init_range, self._init_range)
+        self.embed_dropout = _EmbeddingDropout(self.embed_tokens, config.embed_dropout)
+
+        # LSTM layers
+        self.layers = nn.ModuleList()
+        for i in range(config.num_hidden_layers):
+            input_size = config.hidden_size if i == 0 else config.intermediate_size
+            output_size = (config.intermediate_size if i != config.num_hidden_layers - 1
+                          else config.hidden_size) // self.num_directions
+            lstm = nn.LSTM(input_size, output_size, num_layers=1,
+                          batch_first=True, bidirectional=config.bidirectional)
+            self.layers.append(_WeightDropout(lstm, config.weight_dropout))
+
+        # Dropout
+        self.input_dropout = _RNNDropout(config.input_dropout)
+        self.hidden_dropout = nn.ModuleList([
+            _RNNDropout(config.hidden_dropout) for _ in range(config.num_hidden_layers)
+        ])
+
+        self._hidden_state = None
+        self.reset()
+
+    def reset(self):
+        """Reset LSTM hidden states."""
+        self._hidden_state = [self._init_hidden(i) for i in range(self.num_layers)]
+
+    def _init_hidden(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        nh = (self.intermediate_size if layer_idx != self.num_layers - 1
+              else self.hidden_size) // self.num_directions
+        weight = next(self.parameters())
+        return (weight.new_zeros(self.num_directions, self._batch_size, nh),
+                weight.new_zeros(self.num_directions, self._batch_size, nh))
+
+    def _resize_hidden(self, batch_size: int):
+        new_hidden = []
+        for i in range(self.num_layers):
+            nh = (self.intermediate_size if i != self.num_layers - 1
+                  else self.hidden_size) // self.num_directions
+            h, c = self._hidden_state[i]
+
+            if self._batch_size < batch_size:
+                h = torch.cat([h, h.new_zeros(self.num_directions, batch_size - self._batch_size, nh)], dim=1)
+                c = torch.cat([c, c.new_zeros(self.num_directions, batch_size - self._batch_size, nh)], dim=1)
+            elif self._batch_size > batch_size:
+                h = h[:, :batch_size].contiguous()
+                c = c[:, :batch_size].contiguous()
+            new_hidden.append((h, c))
+
+        self._hidden_state = new_hidden
+        self._batch_size = batch_size
+
+    def forward(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """Returns hidden states for all positions: (batch, seq_len, hidden_size)"""
+        batch_size, seq_len = input_ids.shape
+
+        if batch_size != self._batch_size:
+            self._resize_hidden(batch_size)
+
+        hidden = self.input_dropout(self.embed_dropout(input_ids))
+
+        new_hidden = []
+        for i, (layer, hdp) in enumerate(zip(self.layers, self.hidden_dropout)):
+            hidden, h = layer(hidden, self._hidden_state[i])
+            new_hidden.append(h)
+            if i != self.num_layers - 1:
+                hidden = hdp(hidden)
+
+        self._hidden_state = [(h.detach(), c.detach()) for h, c in new_hidden]
+        return hidden
+
+
+class _LMHead(nn.Module):
+    """Language modeling head."""
+
+    _init_range = 0.1
+
+    def __init__(self, config: LookingGlassConfig, embed_tokens: Optional[nn.Embedding] = None):
+        super().__init__()
+        self.output_dropout = _RNNDropout(config.output_dropout)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=config.output_bias)
+        self.decoder.weight.data.uniform_(-self._init_range, self._init_range)
+
+        if config.output_bias:
+            self.decoder.bias.data.zero_()
+
+        if embed_tokens is not None and config.tie_weights:
+            self.decoder.weight = embed_tokens.weight
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.output_dropout(hidden_states))
+
+
+# =============================================================================
+# Models
+# =============================================================================
+
+class LookingGlass(nn.Module):
+    """
+    LookingGlass encoder model.
+
+    Outputs sequence embeddings for downstream tasks (classification, clustering, etc.).
+    Uses last-token embedding by default, matching original LookingGlass.
+
+    Example:
+        >>> model = LookingGlass.from_pretrained('lookingglass-v1')
+        >>> tokenizer = LookingGlassTokenizer()
+        >>> inputs = tokenizer("GATTACA", return_tensors=True)
+        >>> embeddings = model.get_embeddings(inputs['input_ids'])  # (1, 104)
+    """
+
+    config_class = LookingGlassConfig
+
+    def __init__(self, config: Optional[LookingGlassConfig] = None):
+        super().__init__()
+        self.config = config or LookingGlassConfig()
+        self.encoder = _AWDLSTMEncoder(self.config)
+
+    def reset(self):
+        """Reset hidden states."""
+        self.encoder.reset()
+
+    def forward(self, input_ids: torch.LongTensor, **kwargs) -> torch.Tensor:
+        """
+        Forward pass. Returns last-token embeddings.
+
+        Args:
+            input_ids: Token indices (batch, seq_len)
+
+        Returns:
+            Embeddings (batch, hidden_size)
+        """
+        return self.get_embeddings(input_ids)
+
+    def get_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Get sequence embeddings using last-token pooling (original LG method).
+
+        Resets hidden state before encoding for deterministic results.
+
+        Args:
+            input_ids: Token indices (batch, seq_len)
+
+        Returns:
+            Embeddings (batch, hidden_size)
+        """
+        self.encoder.reset()
+        hidden = self.encoder(input_ids)  # (batch, seq_len, hidden_size)
+        return hidden[:, -1]  # last token
+
+    def get_hidden_states(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Get hidden states for all positions.
+
+        Resets hidden state before encoding for deterministic results.
+
+        Args:
+            input_ids: Token indices (batch, seq_len)
+
+        Returns:
+            Hidden states (batch, seq_len, hidden_size)
+        """
+        self.encoder.reset()
+        return self.encoder(input_ids)
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+        self.config.save_pretrained(save_directory)
+        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_path: str, config: Optional[LookingGlassConfig] = None) -> "LookingGlass":
+        config = config or LookingGlassConfig.from_pretrained(pretrained_path)
+        model = cls(config)
+
+        if _is_hf_hub_id(pretrained_path):
+            model_path = _download_from_hub(pretrained_path, "pytorch_model.bin")
+        else:
+            model_path = os.path.join(pretrained_path, "pytorch_model.bin")
+
+        if os.path.exists(model_path):
+            state_dict = torch.load(model_path, map_location='cpu')
+            # Only load encoder weights
+            encoder_state_dict = {k: v for k, v in state_dict.items()
+                                  if not k.startswith('lm_head.')}
+            model.load_state_dict(encoder_state_dict, strict=False)
+
+        return model
+
+
+class LookingGlassLM(nn.Module):
+    """
+    LookingGlass with language modeling head.
+
+    Full model for next-token prediction. Can also extract embeddings.
+
+    Example:
+        >>> model = LookingGlassLM.from_pretrained('lookingglass-v1')
+        >>> tokenizer = LookingGlassTokenizer()
+        >>> inputs = tokenizer("GATTACA", return_tensors=True)
+        >>> logits = model(inputs['input_ids'])  # (1, 8, 8)
+        >>> embeddings = model.get_embeddings(inputs['input_ids'])  # (1, 104)
+    """
+
+    config_class = LookingGlassConfig
+
+    def __init__(self, config: Optional[LookingGlassConfig] = None):
+        super().__init__()
+        self.config = config or LookingGlassConfig()
+        self.encoder = _AWDLSTMEncoder(self.config)
+        self.lm_head = _LMHead(
+            self.config,
+            embed_tokens=self.encoder.embed_tokens if self.config.tie_weights else None
+        )
+
+    def reset(self):
+        """Reset hidden states."""
+        self.encoder.reset()
+
+    def forward(self, input_ids: torch.LongTensor, **kwargs) -> torch.Tensor:
+        """
+        Forward pass. Returns logits for next-token prediction.
+
+        Args:
+            input_ids: Token indices (batch, seq_len)
+
+        Returns:
+            Logits (batch, seq_len, vocab_size)
+        """
+        hidden = self.encoder(input_ids)
+        return self.lm_head(hidden)
+
+    def get_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Get sequence embeddings using last-token pooling.
+
+        Resets hidden state before encoding for deterministic results.
+
+        Args:
+            input_ids: Token indices (batch, seq_len)
+
+        Returns:
+            Embeddings (batch, hidden_size)
+        """
+        self.encoder.reset()
+        hidden = self.encoder(input_ids)
+        return hidden[:, -1]
+
+    def get_hidden_states(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Get hidden states for all positions.
+
+        Resets hidden state before encoding for deterministic results.
+
+        Args:
+            input_ids: Token indices (batch, seq_len)
+
+        Returns:
+            Hidden states (batch, seq_len, hidden_size)
+        """
+        self.encoder.reset()
+        return self.encoder(input_ids)
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+        self.config.save_pretrained(save_directory)
+        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_path: str, config: Optional[LookingGlassConfig] = None) -> "LookingGlassLM":
+        config = config or LookingGlassConfig.from_pretrained(pretrained_path)
+        model = cls(config)
+
+        if _is_hf_hub_id(pretrained_path):
+            model_path = _download_from_hub(pretrained_path, "pytorch_model.bin")
+        else:
+            model_path = os.path.join(pretrained_path, "pytorch_model.bin")
+
+        if os.path.exists(model_path):
+            state_dict = torch.load(model_path, map_location='cpu')
+            model.load_state_dict(state_dict, strict=False)
+
+        return model
+
+
+# =============================================================================
+# Weight Loading
+# =============================================================================
+
+def load_original_weights(model: Union[LookingGlass, LookingGlassLM], weights_path: str) -> None:
+    """
+    Load weights from original fastai-trained LookingGlass checkpoint.
+
+    Args:
+        model: Model to load weights into
+        weights_path: Path to LookingGlass.pth or LookingGlass_enc.pth
+    """
+    checkpoint = torch.load(weights_path, map_location='cpu')
+
+    if 'model' in checkpoint:
+        state_dict = checkpoint['model']
+    else:
+        state_dict = checkpoint
+
+    is_lm_model = isinstance(model, LookingGlassLM)
+
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        if '.module.weight_hh_l0' in k:
+            continue
+
+        if k.startswith('0.'):
+            new_k = k[2:]
+            new_k = new_k.replace('encoder.', 'embed_tokens.')
+            new_k = new_k.replace('encoder_dp.emb.', 'embed_tokens.')
+            new_k = new_k.replace('rnns.', 'layers.')
+            new_k = new_k.replace('hidden_dps.', 'hidden_dropout.')
+            new_k = new_k.replace('input_dp.', 'input_dropout.')
+            new_state_dict['encoder.' + new_k] = v
+
+        elif k.startswith('1.') and is_lm_model:
+            new_k = k[2:]
+            new_k = new_k.replace('output_dp.', 'output_dropout.')
+            new_state_dict['lm_head.' + new_k] = v
+
+        else:
+            new_k = k.replace('encoder.', 'embed_tokens.')
+            new_k = new_k.replace('encoder_dp.emb.', 'embed_tokens.')
+            new_k = new_k.replace('rnns.', 'layers.')
+            new_k = new_k.replace('hidden_dps.', 'hidden_dropout.')
+            new_k = new_k.replace('input_dp.', 'input_dropout.')
+            new_state_dict['encoder.' + new_k] = v
+
+    model.load_state_dict(new_state_dict, strict=False)
+
+
+def convert_checkpoint(input_path: str, output_dir: str) -> None:
+    """Convert original checkpoint to new format."""
+    config = LookingGlassConfig()
+    model = LookingGlassLM(config)
+    load_original_weights(model, input_path)
+    model.save_pretrained(output_dir)
+
+    tokenizer = LookingGlassTokenizer()
+    tokenizer.save_pretrained(output_dir)
+    print(f"Saved to {output_dir}")
+
+
+# =============================================================================
+# CLI
+# =============================================================================
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser(description='LookingGlass DNA Language Model')
+    parser.add_argument('--convert', type=str, help='Convert original weights')
+    parser.add_argument('--output', type=str, default='./lookingglass-v1', help='Output directory')
+    parser.add_argument('--test', action='store_true', help='Run tests')
+    args = parser.parse_args()
+
+    if args.convert:
+        convert_checkpoint(args.convert, args.output)
+
+    elif args.test:
+        print("Testing LookingGlass...\n")
+
+        tokenizer = LookingGlassTokenizer()
+        print(f"Vocab: {tokenizer.vocab}")
+        print(f"BOS token added: {tokenizer.add_bos_token}")
+        print(f"EOS token added: {tokenizer.add_eos_token}")
+
+        inputs = tokenizer("GATTACA", return_tensors=True)
+        print(f"\nTokenized 'GATTACA': {inputs['input_ids']}")
+        print(f"Decoded: {tokenizer.decode(inputs['input_ids'][0])}")
+
+        config = LookingGlassConfig()
+        print(f"\nConfig: bidirectional={config.bidirectional}")
+
+        # Test LookingGlass (encoder)
+        encoder = LookingGlass(config)
+        print(f"\nLookingGlass params: {sum(p.numel() for p in encoder.parameters()):,}")
+
+        encoder.eval()
+        with torch.no_grad():
+            emb = encoder.get_embeddings(inputs['input_ids'])
+            print(f"Embeddings shape: {emb.shape}")
+
+        # Test LookingGlassLM
+        lm = LookingGlassLM(config)
+        print(f"\nLookingGlassLM params: {sum(p.numel() for p in lm.parameters()):,}")
+
+        lm.eval()
+        with torch.no_grad():
+            logits = lm(inputs['input_ids'])
+            emb = lm.get_embeddings(inputs['input_ids'])
+            print(f"Logits shape: {logits.shape}")
+            print(f"Embeddings shape: {emb.shape}")
+
+        print("\nAll tests passed!")
+
+    else:
+        parser.print_help()

lookingglass_classifier.py

@@ -0,0 +1,258 @@
+"""
+LookingGlass Classifiers - Fine-tuned DNA sequence classifiers
+
+Pure PyTorch implementation of LookingGlass classifiers from the paper.
+Uses LookingGlass encoder with classification head.
+
+Usage:
+    from lookingglass_classifier import LookingGlassClassifier, LookingGlassTokenizer
+
+    model = LookingGlassClassifier.from_pretrained('.')
+    tokenizer = LookingGlassTokenizer()
+
+    inputs = tokenizer(["GATTACA"], return_tensors=True)
+    logits = model(inputs['input_ids'])  # (batch, num_classes)
+    predictions = logits.argmax(dim=-1)
+"""
+
+import json
+import os
+from dataclasses import dataclass, asdict, field
+from typing import Optional, List
+
+import torch
+import torch.nn as nn
+
+from lookingglass import (
+    LookingGlassConfig,
+    LookingGlassTokenizer,
+    _AWDLSTMEncoder,
+    _is_hf_hub_id,
+    _download_from_hub,
+)
+
+__version__ = "1.1.0"
+__all__ = ["LookingGlassClassifierConfig", "LookingGlassClassifier", "LookingGlassTokenizer"]
+
+
+@dataclass
+class LookingGlassClassifierConfig(LookingGlassConfig):
+    """Configuration for LookingGlass classifier."""
+    num_classes: int = 2
+    classifier_hidden: int = 50
+    classifier_dropout: float = 0.0
+    class_names: List[str] = field(default_factory=list)
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+        with open(os.path.join(save_directory, "config.json"), 'w') as f:
+            json.dump(self.to_dict(), f, indent=2)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_path: str) -> "LookingGlassClassifierConfig":
+        if _is_hf_hub_id(pretrained_path):
+            try:
+                config_path = _download_from_hub(pretrained_path, "config.json")
+            except Exception:
+                return cls()
+        elif os.path.isdir(pretrained_path):
+            config_path = os.path.join(pretrained_path, "config.json")
+        else:
+            config_path = pretrained_path
+
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                config_dict = json.load(f)
+            valid_fields = {f.name for f in cls.__dataclass_fields__.values()}
+            return cls(**{k: v for k, v in config_dict.items() if k in valid_fields})
+        return cls()
+
+
+class LookingGlassClassifier(nn.Module):
+    """
+    LookingGlass with classification head.
+
+    Uses concat pooling (max + mean + last) followed by classification layers.
+
+    Example:
+        >>> model = LookingGlassClassifier.from_pretrained('.')
+        >>> tokenizer = LookingGlassTokenizer()
+        >>> inputs = tokenizer("GATTACA", return_tensors=True)
+        >>> logits = model(inputs['input_ids'])  # (1, num_classes)
+        >>> prediction = logits.argmax(dim=-1)
+    """
+
+    def __init__(self, config: Optional[LookingGlassClassifierConfig] = None):
+        super().__init__()
+        self.config = config or LookingGlassClassifierConfig()
+        self.encoder = _AWDLSTMEncoder(self.config)
+
+        # Concat pooling: max + mean + last = 3 * hidden_size
+        pooled_size = 3 * self.config.hidden_size
+
+        # Classification head: BatchNorm -> Linear -> ReLU -> BatchNorm -> Linear
+        self.classifier = nn.Sequential(
+            nn.BatchNorm1d(pooled_size),
+            nn.Dropout(self.config.classifier_dropout),
+            nn.Linear(pooled_size, self.config.classifier_hidden),
+            nn.ReLU(),
+            nn.BatchNorm1d(self.config.classifier_hidden),
+            nn.Dropout(self.config.classifier_dropout),
+            nn.Linear(self.config.classifier_hidden, self.config.num_classes),
+        )
+
+    def forward(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """
+        Forward pass returning classification logits.
+
+        Args:
+            input_ids: Token indices (batch, seq_len)
+
+        Returns:
+            Logits (batch, num_classes)
+        """
+        self.encoder.reset()
+        hidden = self.encoder(input_ids)  # (batch, seq_len, hidden_size)
+
+        # Concat pooling: max, mean, last
+        max_pool = hidden.max(dim=1).values
+        mean_pool = hidden.mean(dim=1)
+        last_pool = hidden[:, -1]
+        pooled = torch.cat([max_pool, mean_pool, last_pool], dim=-1)
+
+        return self.classifier(pooled)
+
+    def predict(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """Return predicted class indices."""
+        logits = self.forward(input_ids)
+        return logits.argmax(dim=-1)
+
+    def predict_proba(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """Return class probabilities."""
+        logits = self.forward(input_ids)
+        return torch.softmax(logits, dim=-1)
+
+    def get_embeddings(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        """Get sequence embeddings (last token) from encoder."""
+        self.encoder.reset()
+        hidden = self.encoder(input_ids)
+        return hidden[:, -1]
+
+    def save_pretrained(self, save_directory: str):
+        os.makedirs(save_directory, exist_ok=True)
+        self.config.save_pretrained(save_directory)
+        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_path: str, config: Optional[LookingGlassClassifierConfig] = None
+    ) -> "LookingGlassClassifier":
+        config = config or LookingGlassClassifierConfig.from_pretrained(pretrained_path)
+        model = cls(config)
+
+        if _is_hf_hub_id(pretrained_path):
+            model_path = _download_from_hub(pretrained_path, "pytorch_model.bin")
+        else:
+            model_path = os.path.join(pretrained_path, "pytorch_model.bin")
+
+        if os.path.exists(model_path):
+            state_dict = torch.load(model_path, map_location='cpu')
+            model.load_state_dict(state_dict, strict=False)
+
+        return model
+
+
+def convert_classifier_weights(
+    original_path: str,
+    output_dir: str,
+    num_classes: int,
+    class_names: Optional[List[str]] = None,
+):
+    """
+    Convert original fastai classifier weights to pure PyTorch format.
+
+    Args:
+        original_path: Path to original .pth file
+        output_dir: Output directory for converted model
+        num_classes: Number of output classes
+        class_names: Optional list of class names
+    """
+    print(f"Loading weights from {original_path}...")
+    original = torch.load(original_path, map_location='cpu')
+    if 'model' in original:
+        original = original['model']
+
+    # Create config
+    config = LookingGlassClassifierConfig(
+        num_classes=num_classes,
+        classifier_hidden=50,
+        class_names=class_names or [],
+    )
+
+    # Create model
+    model = LookingGlassClassifier(config)
+
+    # Map weights
+    new_state = {}
+
+    # Encoder weights
+    weight_map = {
+        '0.module.encoder.weight': 'encoder.embed_tokens.weight',
+        '0.module.encoder_dp.emb.weight': 'encoder.embed_dropout.embedding.weight',
+    }
+
+    for i in range(3):
+        weight_map.update({
+            f'0.module.rnns.{i}.weight_hh_l0_raw': f'encoder.layers.{i}.weight_hh_l0_raw',
+            f'0.module.rnns.{i}.module.weight_ih_l0': f'encoder.layers.{i}.module.weight_ih_l0',
+            f'0.module.rnns.{i}.module.weight_hh_l0': f'encoder.layers.{i}.module.weight_hh_l0',
+            f'0.module.rnns.{i}.module.bias_ih_l0': f'encoder.layers.{i}.module.bias_ih_l0',
+            f'0.module.rnns.{i}.module.bias_hh_l0': f'encoder.layers.{i}.module.bias_hh_l0',
+        })
+
+    # Classifier head weights
+    # Original: 1.layers.{0,2,4,6} -> our Sequential indices
+    classifier_map = {
+        '1.layers.0.weight': 'classifier.0.weight',
+        '1.layers.0.bias': 'classifier.0.bias',
+        '1.layers.0.running_mean': 'classifier.0.running_mean',
+        '1.layers.0.running_var': 'classifier.0.running_var',
+        '1.layers.0.num_batches_tracked': 'classifier.0.num_batches_tracked',
+        '1.layers.2.weight': 'classifier.2.weight',
+        '1.layers.2.bias': 'classifier.2.bias',
+        '1.layers.4.weight': 'classifier.4.weight',
+        '1.layers.4.bias': 'classifier.4.bias',
+        '1.layers.4.running_mean': 'classifier.4.running_mean',
+        '1.layers.4.running_var': 'classifier.4.running_var',
+        '1.layers.4.num_batches_tracked': 'classifier.4.num_batches_tracked',
+        '1.layers.6.weight': 'classifier.6.weight',
+        '1.layers.6.bias': 'classifier.6.bias',
+    }
+    weight_map.update(classifier_map)
+
+    for old_key, new_key in weight_map.items():
+        if old_key in original:
+            new_state[new_key] = original[old_key]
+
+    # Load and save
+    model.load_state_dict(new_state, strict=False)
+
+    os.makedirs(output_dir, exist_ok=True)
+    config.save_pretrained(output_dir)
+    torch.save(model.state_dict(), os.path.join(output_dir, "pytorch_model.bin"))
+
+    print(f"Saved to {output_dir}")
+    return model
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Convert LookingGlass classifier weights")
+    parser.add_argument("--input", required=True, help="Path to original .pth file")
+    parser.add_argument("--output", required=True, help="Output directory")
+    parser.add_argument("--num-classes", type=int, required=True, help="Number of classes")
+    parser.add_argument("--class-names", nargs="+", help="Class names")
+
+    args = parser.parse_args()
+    convert_classifier_weights(args.input, args.output, args.num_classes, args.class_names)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6104341f1668118250861bdcc9dd17b26af44de9b84e27e2a1bb3ee0f74f67fb
+size 68121438