returnn 1.20241209.230715__tar.gz → 1.20241210.150216__tar.gz

{returnn-1.20241209.230715 → returnn-1.20241210.150216}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20241209.230715
+Version: 1.20241210.150216
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20241210.150216/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241210.150216'
+long_version = '1.20241210.150216+git.d0d1801'

{returnn-1.20241209.230715 → returnn-1.20241210.150216}/returnn/datasets/audio.py

@@ -167,6 +167,7 @@ class OggZipDataset(CachedDataset2):
         else:
             assert isinstance(epoch_wise_filter, EpochWiseFilter)
             self.epoch_wise_filter = epoch_wise_filter
+        self._seq_index_by_tag: Optional[Dict[str, int]] = None
         self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
 
     def _read(self, filename, zip_index):
@@ -288,6 +289,11 @@ class OggZipDataset(CachedDataset2):
             segment_file_handle = open(segment_file)
             self.segments = set(segment_file_handle.read().splitlines())
 
+    def _lazy_init_seq_index_by_tag(self):
+        if self._seq_index_by_tag is not None:
+            return
+        self._seq_index_by_tag = {self._get_tag_from_info_dict(seq): i for i, seq in enumerate(self._data)}
+
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
         """
         If random_shuffle_epoch1, for epoch 1 with "random" ordering, we leave the given order as is.
@@ -327,12 +333,12 @@ class OggZipDataset(CachedDataset2):
         if seq_order is not None:
             self._seq_order = seq_order
         elif seq_list is not None:
-            seqs = {self._get_tag_from_info_dict(seq): i for i, seq in enumerate(self._data)}
+            self._lazy_init_seq_index_by_tag()
             for seq_tag in seq_list:
-                assert seq_tag in seqs, "Did not find all requested seqs. We have eg: %s" % (
+                assert seq_tag in self._seq_index_by_tag, "Did not find all requested seqs. We have eg: %s" % (
                     self._get_tag_from_info_dict(self._data[0]),
                 )
-            self._seq_order = [seqs[seq_tag] for seq_tag in seq_list]
+            self._seq_order = [self._seq_index_by_tag[seq_tag] for seq_tag in seq_list]
         else:
             num_seqs = len(self._data)
             self._seq_order = self.get_seq_order_for_epoch(epoch=epoch, num_seqs=num_seqs, get_seq_len=get_seq_len)

{returnn-1.20241209.230715 → returnn-1.20241210.150216}/returnn/datasets/basic.py

@@ -330,6 +330,8 @@ class Dataset:
             assert chunk_step.max_value() > 0, "chunking step must be positive (for some key)"
         return chunk_size, chunk_step, None
 
+    _seq_list_file_cache: Dict[str, Union[List[str], Dict[str, List[str]]]] = {}
+
     @staticmethod
     def _load_seq_list_file(
         filename: str, *, use_cache_manager: bool = False, expect_list: bool = True
@@ -340,6 +342,12 @@ class Dataset:
         :param expect_list:
         :return: seq list file content. usually a list of seqs.
         """
+        if filename in Dataset._seq_list_file_cache:
+            seq_list = Dataset._seq_list_file_cache[filename]
+            if expect_list:
+                assert isinstance(seq_list, list)
+            return seq_list
+        _filename = filename
         if use_cache_manager:
             import returnn.util.basic
 
@@ -367,6 +375,7 @@ class Dataset:
                 assert isinstance(seq_list, list)
         else:
             seq_list = _open("rt").read().splitlines()
+        Dataset._seq_list_file_cache[_filename] = seq_list
         return seq_list
 
     def _sliding_window(self, xr):

{returnn-1.20241209.230715 → returnn-1.20241210.150216}/returnn/datasets/lm.py

@@ -46,8 +46,10 @@ class LmDataset(CachedDataset2):
     def __init__(
         self,
         corpus_file,
+        *,
         use_cache_manager=False,
         skip_empty_lines=True,
+        seq_list_file=None,
         orth_vocab=None,
         orth_symbols_file=None,
         orth_symbols_map_file=None,
@@ -89,6 +91,8 @@ class LmDataset(CachedDataset2):
         :param str|()->str|list[str]|()->list[str] corpus_file: Bliss XML or line-based txt. optionally can be gzip.
         :param bool use_cache_manager: uses :func:`returnn.util.basic.cf`
         :param bool skip_empty_lines: for line-based txt
+        :param str|list[str]|None seq_list_file: optional custom seq tags to use instead of the "line-%i" seq tags.
+            Pickle (.pkl) or txt (line-based seq tags). Optionally gzipped (.gz).
         :param dict[str,typing.Any]|Vocabulary orth_vocab:
         :param str|()->str|None orth_symbols_file: a text file containing a list of orthography symbols
         :param str|()->str|None orth_symbols_map_file: either a list of orth symbols, each line: "<symbol> <index>",
@@ -119,6 +123,7 @@ class LmDataset(CachedDataset2):
         self._corpus_file = corpus_file
         self._use_cache_manager = use_cache_manager
         self._skip_empty_lines = skip_empty_lines
+        self._seq_list_file = seq_list_file
         self._orth_symbols_file = orth_symbols_file
         self._orth_symbols_map_file = orth_symbols_map_file
         self._orth_replace_map_file = orth_replace_map_file
@@ -274,6 +279,8 @@ class LmDataset(CachedDataset2):
         self._orth_files: Optional[List[BinaryIO]] = None
         self._orth_mmaps = None
         self._orths_offsets_and_lens: Optional[List[Tuple[int, int]]] = None  # will be loaded in _lazy_init
+        self._seq_list: Optional[List[str]] = None
+        self._seq_index_by_tag: Optional[dict[str, int]] = None
 
         self.next_orth_idx = 0
         self.next_seq_idx = 0
@@ -301,6 +308,7 @@ class LmDataset(CachedDataset2):
         total_bytes_read = 0
         orths = []
         self._orths_offsets_and_lens = orths
+        lens_per_corpus_file = []
         start_time = time.time()
         last_print_time = start_time
 
@@ -341,6 +349,7 @@ class LmDataset(CachedDataset2):
         # If a list of files is provided, concatenate all.
         if not isinstance(corpus_file, list):
             corpus_file = [corpus_file]
+        prev_orth_len = 0
         for file_name in corpus_file:
             if self._use_cache_manager:
                 file_name = cf(file_name)
@@ -376,10 +385,31 @@ class LmDataset(CachedDataset2):
                     pos = next_new_line + 1
                     _maybe_report_status()
 
+            lens_per_corpus_file.append(len(orths) - prev_orth_len)
+            prev_orth_len = len(orths)
+
         if tmp_file is not None:
             tmp_file.flush()
             self._orth_mmaps[tmp_file_orth_files_index] = mmap.mmap(tmp_file.fileno(), 0, flags=mmap.MAP_PRIVATE)
 
+        if self._seq_list_file:
+            if isinstance(self._seq_list_file, str):
+                seq_list: List[str] = self._load_seq_list_file(
+                    self._seq_list_file, use_cache_manager=self._use_cache_manager
+                )
+            elif isinstance(self._seq_list_file, list):
+                assert len(self._seq_list_file) == len(lens_per_corpus_file)
+                seq_list: List[str] = []
+                for i, fn in enumerate(self._seq_list_file):
+                    seq_list_: List[str] = self._load_seq_list_file(fn, use_cache_manager=self._use_cache_manager)
+                    assert len(seq_list_) == lens_per_corpus_file[i]
+                    seq_list.extend(seq_list_)
+            else:
+                raise TypeError(f"invalid seq_list_file type {type(self._seq_list_file).__name__}")
+            assert isinstance(seq_list, list)
+            assert len(self._orths_offsets_and_lens) == len(seq_list)
+            self._seq_list = seq_list
+
         print(
             f"  done, loaded {len(self._orths_offsets_and_lens)} sequences,"
             f" {human_bytes_size(total_bytes_read)},"
@@ -445,8 +475,16 @@ class LmDataset(CachedDataset2):
         if seq_order is not None:
             self.seq_order = seq_order
         elif seq_list is not None:
-            assert all(s.startswith(self._tag_prefix) for s in seq_list)
-            self.seq_order = [int(s[len(self._tag_prefix) :]) for s in seq_list]
+            # Might not be initialized. Can even do without init. Thus check seq_list_file.
+            if self._seq_list_file is None:
+                assert all(s.startswith(self._tag_prefix) for s in seq_list)
+                self.seq_order = [int(s[len(self._tag_prefix) :]) for s in seq_list]
+            else:
+                # Need seq list for this. Just do the lazy init now.
+                self._lazy_init()
+                if self._seq_index_by_tag is None:
+                    self._seq_index_by_tag = {tag: i for (i, tag) in enumerate(self._seq_list)}
+                self.seq_order = [self._seq_index_by_tag[s] for s in seq_list]
         elif epoch is None:
             self.seq_order = []
         else:
@@ -486,6 +524,9 @@ class LmDataset(CachedDataset2):
 
     def get_all_tags(self) -> List[str]:
         """:return: all seq tags"""
+        self._lazy_init()
+        if self._seq_list is not None:
+            return self._seq_list
         num_seqs = self.get_total_num_seqs()
         return [self._tag_prefix + str(line_nr) for line_nr in range(num_seqs)]
 
@@ -525,10 +566,11 @@ class LmDataset(CachedDataset2):
             # get sequence for the next index given by seq_order
             idx, offset, len_ = self._orths_offsets_and_lens[true_idx]
             orth = self._orth_mmaps[idx][offset : offset + len_].decode("utf8").strip()
-            seq_tag = self._tag_prefix + str(true_idx)
+            if self._seq_list is None:
+                seq_tag = self._tag_prefix + str(true_idx)
+            else:
+                seq_tag = self._seq_list[true_idx]
             self.next_orth_idx += 1
-            if orth == "</s>":
-                continue  # special sentence end symbol. empty seq, ignore.
 
             if self.orth_vocab is not None:
                 data = numpy.array(self.orth_vocab.get_seq(orth), dtype=self.dtype)

{returnn-1.20241209.230715 → returnn-1.20241210.150216}/returnn/frontend/attention.py

@@ -325,7 +325,7 @@ class RotaryPosCausalSelfAttention(CausalSelfAttention):
         q = _apply_rope(
             q,
             (
-                rf.gather(pos_enc, axis=hist_dim, indices=hist_dim.dyn_size_ext - 1)
+                rf.gather(pos_enc, axis=hist_dim, indices=hist_dim.get_size_tensor() - 1)
                 if axis == single_step_dim
                 else rf.replace_dim(pos_enc, in_dim=hist_dim, out_dim=axis)[0]
             ),
@@ -503,6 +503,8 @@ def _rel_pos_enc_shift(x: Tensor, axis: Dim, pos_emb_spatial_dim: Dim, hist_dim:
     :param hist_dim: T' (equal to T but separate dim)
     :return: [B,H,T,T']
     """
+    if pos_emb_spatial_dim == hist_dim:  # happens for single_step_dim
+        return x  # no shift needed
     batch_dims = x.remaining_dims((axis, pos_emb_spatial_dim))
     x_padded, (pos_emb_spatial_dim_,) = rf.pad(
         x, axes=[pos_emb_spatial_dim], padding=[(1, 0)], value=0.0
@@ -604,6 +606,7 @@ class RelPosCausalSelfAttention(CausalSelfAttention):
             pos_emb, pos_emb_spatial_dim = relative_positional_encoding(
                 query_spatial_dim=axis, key_value_spatial_dim=hist_dim, feat_dim=self.pos_emb_feat_dim
             )
+        # pos_emb_spatial_dim is 2*time1-1 if axis!=single_step_dim, else time1
         if self.pos_emb_dropout:
             pos_emb = rf.dropout(pos_emb, self.pos_emb_dropout)
         if self.linear_pos is not None:
@@ -850,7 +853,8 @@ def _make_indices(
     if query_spatial_dim == single_step_dim:
         indices = kv_pos_vec
         out_spatial_dim = key_value_spatial_dim
-        assert query_offset is None  # not sure if any custom query offset makes sense?
+        # not sure if any custom query offset makes sense?
+        assert query_offset is None or (isinstance(query_offset, int) and query_offset == 0)
         # Assume the kv are the accumulated history, and query is cur frame of it,
         # corresponding to the last frame of the kv.
         query_offset = key_value_spatial_dim.get_size_tensor() - 1

{returnn-1.20241209.230715 → returnn-1.20241210.150216}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20241209.230715
+Version: 1.20241210.150216
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20241209.230715 → returnn-1.20241210.150216}/tests/test_rf_attention.py

@@ -3,7 +3,7 @@ RETURNN frontend (returnn.frontend) tests
 """
 
 from __future__ import annotations
-from typing import Tuple
+from typing import Union, Tuple
 import numpy as np
 import numpy.testing
 import _setup_test_env  # noqa
@@ -440,6 +440,84 @@ def test_rope_causal_self_att():
     print("  all matched!")
 
 
+def test_causal_self_att_variants_single_step_vs_full_seq():
+    from returnn.tensor import single_step_dim
+
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7 * 2, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, model: Union[rf.CausalSelfAttention], extern_data: TensorDict):
+        x = extern_data["data"]
+
+        out_seq_level, _ = model(x, axis=time_dim)
+        out_seq_level.mark_as_output("out_seq_level", shape=[batch_dim, time_dim, model.out_dim])
+
+        out_seq_level_explicit_initial_state, _ = model(
+            x, axis=time_dim, state=model.default_initial_state(batch_dims=[batch_dim])
+        )
+        out_seq_level_explicit_initial_state.mark_as_output(
+            "out_seq_level_explicit_initial_state", shape=[batch_dim, time_dim, model.out_dim]
+        )
+
+        def _body(
+            _x: Tensor, _state: Union[rf.CausalSelfAttentionState]
+        ) -> Tuple[Tensor, Union[rf.CausalSelfAttentionState]]:
+            return model(_x, axis=single_step_dim, state=_state)
+
+        out_single_steps, _, _ = rf.scan(
+            spatial_dim=time_dim,
+            xs=x,
+            body=_body,
+            ys=Tensor("y", dims=[batch_dim, model.out_dim], dtype="float32"),
+            initial=model.default_initial_state(batch_dims=[batch_dim]),
+        )
+        out_single_steps.mark_as_output("out_single_steps", shape=[batch_dim, time_dim, model.out_dim])
+
+    common_opts = dict(
+        in_dim=in_dim,
+        proj_dim=Dim(5, name="out"),
+        key_dim_total=Dim(21 * 2, name="key-dim-total"),
+        value_dim_total=Dim(33, name="value-dim-total"),
+        num_heads=3,
+    )
+
+    def _make_causal_self_att(**_kwargs):
+        return rf.CausalSelfAttention(**common_opts)
+
+    def _make_rope_causal_self_att(**_kwargs):
+        return rf.RotaryPosCausalSelfAttention(**common_opts)
+
+    def _make_rel_pos_causal_self_att(**_kwargs):
+        return rf.RelPosCausalSelfAttention(**common_opts)
+
+    models = [_make_causal_self_att, _make_rope_causal_self_att, _make_rel_pos_causal_self_att]
+
+    for get_model in models:
+        print("> Testing model:", get_model.__name__)
+        res = run_model(
+            extern_data,
+            get_model,
+            _forward_step,
+            # TF needs TensorArray unstack, not implemented yet
+            test_tensorflow=False,
+        )
+
+        # Check that the single-step and the seq-level output are the same.
+        res_seq_level = res.data["out_seq_level"].raw_tensor
+        for key in ["out_seq_level_explicit_initial_state", "out_single_steps"]:
+            res_other = res.data[key].raw_tensor
+            assert res_seq_level.shape == res_other.shape
+            numpy.testing.assert_allclose(
+                res_other, res_seq_level, atol=1e-5, rtol=1e-5, err_msg=f"output {key} differs"
+            )
+
+
 def test_relative_positional_encoding():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(8, name="in")

returnn-1.20241209.230715/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241209.230715'
-long_version = '1.20241209.230715+git.55c6a02'