Anchor-annotator 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

anchor/undo.py

@@ -6,7 +6,7 @@ import typing
 import pynini.lib
 import sqlalchemy
 from montreal_forced_aligner.data import WordType
-from montreal_forced_aligner.db import File, Pronunciation, Speaker, Utterance, Word, bulk_update
+from montreal_forced_aligner.db import File, Pronunciation, Speaker, Utterance, Word
 from PySide6 import QtCore, QtGui
 from sqlalchemy.orm import make_transient
 
@@ -15,6 +15,7 @@ if typing.TYPE_CHECKING:
         CorpusModel,
         DiarizationModel,
         DictionaryTableModel,
+        FileUtterancesModel,
         SpeakerModel,
         TextFilterQuery,
     )
@@ -38,9 +39,8 @@ class CorpusCommand(QtGui.QUndoCommand):
 
     def redo(self) -> None:
         with self.corpus_model.edit_lock:
-            with self.corpus_model.session() as session:
-                self._redo(session)
-                session.commit()
+            self._redo(self.corpus_model.session)
+            self.corpus_model.session.commit()
             # while True:
             #    try:
             #        with self.corpus_model.session.begin_nested():
@@ -53,9 +53,8 @@ class CorpusCommand(QtGui.QUndoCommand):
 
     def undo(self) -> None:
         with self.corpus_model.edit_lock:
-            with self.corpus_model.session() as session:
-                self._undo(session)
-                session.commit()
+            self._undo(self.corpus_model.session)
+            self.corpus_model.session.commit()
             # while True:
             #    try:
             #        with self.corpus_model.session.begin_nested():
@@ -66,6 +65,12 @@ class CorpusCommand(QtGui.QUndoCommand):
         self.update_data()
 
 
+class FileCommand(CorpusCommand):
+    def __init__(self, file_model: FileUtterancesModel):
+        super().__init__(file_model.corpus_model)
+        self.file_model = file_model
+
+
 class DictionaryCommand(QtGui.QUndoCommand):
     def __init__(self, dictionary_model: DictionaryTableModel):
         super().__init__()
@@ -127,9 +132,9 @@ class SpeakerCommand(QtGui.QUndoCommand):
         self.update_data()
 
 
-class DeleteUtteranceCommand(CorpusCommand):
-    def __init__(self, deleted_utterances: list[Utterance], corpus_model: CorpusModel):
-        super().__init__(corpus_model)
+class DeleteUtteranceCommand(FileCommand):
+    def __init__(self, deleted_utterances: list[Utterance], file_model: FileUtterancesModel):
+        super().__init__(file_model)
         self.deleted_utterances = deleted_utterances
         self.resets_tier = True
         self.channels = [
@@ -147,7 +152,6 @@ class DeleteUtteranceCommand(CorpusCommand):
         for i, utt in enumerate(self.deleted_utterances):
             make_transient(utt)
             for x in utt.phone_intervals:
-                x.duration = None
                 make_transient(x)
             for x in utt.word_intervals:
                 make_transient(x)
@@ -158,95 +162,80 @@ class DeleteUtteranceCommand(CorpusCommand):
     def redo(self) -> None:
         super().redo()
         self.corpus_model.delete_table_utterances(self.deleted_utterances)
+        self.file_model.delete_table_utterances(self.deleted_utterances)
         self.corpus_model.changeCommandFired.emit()
 
     def undo(self) -> None:
         super().undo()
         self.corpus_model.add_table_utterances(self.deleted_utterances)
+        self.file_model.add_table_utterances(self.deleted_utterances)
         self.corpus_model.changeCommandFired.emit()
 
 
-class SplitUtteranceCommand(CorpusCommand):
+class SplitUtteranceCommand(FileCommand):
     def __init__(
         self,
-        split_utterances: list[list[Utterance, ...]],
-        corpus_model: CorpusModel,
+        merged_utterance: Utterance,
+        split_utterances: list[Utterance],
+        file_model: FileUtterancesModel,
         update_table: bool = True,
     ):
-        super().__init__(corpus_model)
+        super().__init__(file_model)
+        self.merged_utterance = merged_utterance
         self.split_utterances = split_utterances
         self.resets_tier = True
         self.update_table = update_table
-        self.channels = [
-            x[0].channel if x[0].channel is not None else 0 for x in self.split_utterances
-        ]
+        self.channels = [x.channel if x.channel is not None else 0 for x in self.split_utterances]
         self.setText(
             QtCore.QCoreApplication.translate("SplitUtteranceCommand", "Split utterances")
         )
 
     def _redo(self, session) -> None:
-        for i, splits in enumerate(self.split_utterances):
-            old_utt = splits[0]
-            split_utts = splits[1:]
-            session.delete(old_utt)
-            for u in split_utts:
-                if u.id is not None:
-                    make_transient(u)
-                for x in u.phone_intervals:
-                    x.duration = None
-                    make_transient(x)
-                for x in u.word_intervals:
-                    make_transient(x)
-                if u.channel is None:
-                    u.channel = self.channels[i]
-                u.duration = None
-                u.kaldi_id = None
-                session.add(u)
-
-    def _undo(self, session) -> None:
-        for i, splits in enumerate(self.split_utterances):
-            old_utt = splits[0]
-            split_utts = splits[1:]
-            if old_utt.channel is None:
-                old_utt.channel = self.channels[i]
-            old_utt.duration = None
-            old_utt.kaldi_id = None
-            make_transient(old_utt)
-            for x in old_utt.phone_intervals:
-                x.duration = None
+        session.delete(self.merged_utterance)
+        for u in self.split_utterances:
+            if u.id is not None:
+                make_transient(u)
+            for x in u.phone_intervals:
                 make_transient(x)
-            for x in old_utt.word_intervals:
+            for x in u.word_intervals:
                 make_transient(x)
-            session.add(old_utt)
-            for u in split_utts:
-                session.delete(u)
+            if u.channel is None:
+                u.channel = self.merged_utterance.channel
+            session.add(u)
+
+    def _undo(self, session) -> None:
+        if self.merged_utterance.channel is None:
+            self.merged_utterance.channel = self.split_utterances[0].channel
+        make_transient(self.merged_utterance)
+        for x in self.merged_utterance.phone_intervals:
+            make_transient(x)
+        for x in self.merged_utterance.word_intervals:
+            make_transient(x)
+        session.add(self.merged_utterance)
+        for u in self.split_utterances:
+            session.delete(u)
 
     def redo(self) -> None:
         super().redo()
-        for splits in self.split_utterances:
-            old_utt = splits[0]
-            split_utts = splits[1:]
-            if self.update_table:
-                self.corpus_model.split_table_utterances(old_utt, split_utts)
+        self.corpus_model.split_table_utterances(self.merged_utterance, self.split_utterances)
+        self.file_model.split_table_utterances(self.merged_utterance, self.split_utterances)
         self.corpus_model.changeCommandFired.emit()
 
     def undo(self) -> None:
         super().undo()
-        for splits in self.split_utterances:
-            old_utt = splits[0]
-            split_utts = splits[1:]
-            self.corpus_model.merge_table_utterances(old_utt, split_utts)
+        self.corpus_model.merge_table_utterances(self.merged_utterance, self.split_utterances)
+        self.file_model.merge_table_utterances(self.merged_utterance, self.split_utterances)
         self.corpus_model.changeCommandFired.emit()
 
 
-class MergeUtteranceCommand(CorpusCommand):
+class MergeUtteranceCommand(FileCommand):
     def __init__(
         self,
         unmerged_utterances: list[Utterance],
         merged_utterance: Utterance,
-        corpus_model: CorpusModel,
+        file_model: FileUtterancesModel,
     ):
-        super().__init__(corpus_model)
+        super().__init__(file_model)
         self.unmerged_utterances = unmerged_utterances
         self.merged_utterance = merged_utterance
         self.resets_tier = True
@@ -263,8 +252,6 @@ class MergeUtteranceCommand(CorpusCommand):
         make_transient(self.merged_utterance)
         if self.merged_utterance.channel is None:
             self.merged_utterance.channel = self.channel
-        self.merged_utterance.kaldi_id = None
-        self.merged_utterance.duration = None
         session.add(self.merged_utterance)
 
     def _undo(self, session) -> None:
@@ -273,24 +260,22 @@ class MergeUtteranceCommand(CorpusCommand):
             if old_utt.channel is None:
                 old_utt.channel = self.channel
             for x in old_utt.phone_intervals:
-                x.duration = None
                 make_transient(x)
             for x in old_utt.word_intervals:
                 make_transient(x)
-            old_utt.duration = None
-            old_utt.kaldi_id = None
             session.add(old_utt)
-        # session.refresh(self.merged_utterance)
         session.delete(self.merged_utterance)
 
     def redo(self) -> None:
         super().redo()
         self.corpus_model.merge_table_utterances(self.merged_utterance, self.unmerged_utterances)
+        self.file_model.merge_table_utterances(self.merged_utterance, self.unmerged_utterances)
         self.corpus_model.changeCommandFired.emit()
 
     def undo(self) -> None:
         super().undo()
         self.corpus_model.split_table_utterances(self.merged_utterance, self.unmerged_utterances)
+        self.file_model.split_table_utterances(self.merged_utterance, self.unmerged_utterances)
         self.corpus_model.changeCommandFired.emit()
 
 
@@ -350,11 +335,10 @@ class MergeSpeakersCommand(CorpusCommand):
         ).update({Speaker.modified: True})
 
 
-class CreateUtteranceCommand(CorpusCommand):
-    def __init__(self, new_utterance: Utterance, corpus_model: CorpusModel):
-        super().__init__(corpus_model)
+class CreateUtteranceCommand(FileCommand):
+    def __init__(self, new_utterance: Utterance, file_model: FileUtterancesModel):
+        super().__init__(file_model)
         self.new_utterance = new_utterance
-        self.resets_tier = True
         self.channel = self.new_utterance.channel
         if self.channel is None:
             self.channel = 0
@@ -377,18 +361,22 @@ class CreateUtteranceCommand(CorpusCommand):
     def redo(self) -> None:
         super().redo()
         self.corpus_model.add_table_utterances([self.new_utterance])
+        self.file_model.add_table_utterances([self.new_utterance])
         self.corpus_model.changeCommandFired.emit()
 
     def undo(self) -> None:
         super().undo()
         self.corpus_model.delete_table_utterances([self.new_utterance])
+        self.file_model.delete_table_utterances([self.new_utterance])
         self.corpus_model.changeCommandFired.emit()
 
 
-class UpdateUtteranceTimesCommand(CorpusCommand):
-    def __init__(self, utterance: Utterance, begin: float, end: float, corpus_model: CorpusModel):
-        super().__init__(corpus_model)
-        self.utterance_id = utterance.id
+class UpdateUtteranceTimesCommand(FileCommand):
+    def __init__(
+        self, utterance: Utterance, begin: float, end: float, file_model: FileUtterancesModel
+    ):
+        super().__init__(file_model)
+        self.utterance = utterance
         self.new_begin = begin
         self.old_begin = utterance.begin
         self.new_end = end
@@ -400,37 +388,31 @@ class UpdateUtteranceTimesCommand(CorpusCommand):
         )
 
     def _redo(self, session) -> None:
-        session.query(Utterance).filter(Utterance.id == self.utterance_id).update(
-            {
-                Utterance.begin: self.new_begin,
-                Utterance.end: self.new_end,
-                Utterance.xvector: None,
-                Utterance.ivector: None,
-                Utterance.features: None,
-            }
-        )
+        self.utterance.begin = self.new_begin
+        self.utterance.end = self.new_end
+        self.utterance.xvector = None
+        self.utterance.ivector = None
+        self.utterance.features = None
+        session.merge(self.utterance)
 
     def _undo(self, session) -> None:
-        session.query(Utterance).filter(Utterance.id == self.utterance_id).update(
-            {
-                Utterance.begin: self.old_begin,
-                Utterance.end: self.old_end,
-                Utterance.xvector: None,
-                Utterance.ivector: None,
-                Utterance.features: None,
-            }
-        )
+        self.utterance.begin = self.old_begin
+        self.utterance.end = self.old_end
+        self.utterance.xvector = None
+        self.utterance.ivector = None
+        self.utterance.features = None
+        session.merge(self.utterance)
 
     def update_data(self):
         super().update_data()
         self.corpus_model.changeCommandFired.emit()
-        self.corpus_model.update_utterance_table_row(self.utterance_id)
+        self.corpus_model.update_utterance_table_row(self.utterance)
 
 
-class UpdateUtteranceTextCommand(CorpusCommand):
-    def __init__(self, utterance: Utterance, new_text: str, corpus_model: CorpusModel):
-        super().__init__(corpus_model)
-        self.utterance_id = utterance.id
+class UpdateUtteranceTextCommand(FileCommand):
+    def __init__(self, utterance: Utterance, new_text: str, file_model: FileUtterancesModel):
+        super().__init__(file_model)
+        self.utterance = utterance
         self.speaker_id = utterance.speaker_id
         self.old_text = utterance.text
         self.new_text = new_text
@@ -445,40 +427,28 @@ class UpdateUtteranceTextCommand(CorpusCommand):
         for w in self.new_text.split():
             if not self.corpus_model.dictionary_model.check_word(w, self.speaker_id):
                 oovs.add(w)
-        session.query(Utterance).filter(Utterance.id == self.utterance_id).update(
-            {
-                Utterance.text: self.new_text,
-                Utterance.normalized_text: self.new_text,  # FIXME: Update this
-                Utterance.oovs: " ".join(oovs),
-                Utterance.ignored: not self.new_text,
-            }
-        )
+        self.utterance.text = self.new_text
+        self.utterance.normalized_text = self.new_text  # FIXME: Update this
+        self.utterance.oovs = " ".join(oovs)
+        self.utterance.ignored = not self.new_text
+        session.merge(self.utterance)
 
     def _undo(self, session) -> None:
         oovs = set()
         for w in self.new_text.split():
             if not self.corpus_model.dictionary_model.check_word(w, self.speaker_id):
                 oovs.add(w)
-        session.query(Utterance).filter(Utterance.id == self.utterance_id).update(
-            {
-                Utterance.text: self.old_text,
-                Utterance.oovs: " ".join(oovs),
-                Utterance.ignored: not self.old_text,
-            }
-        )
-
-    def update_data(self):
-        super().update_data()
-        try:
-            self.corpus_model.update_utterance_table_row(self.utterance_id)
-        except KeyError:
-            pass
+        self.utterance.text = self.old_text
+        self.utterance.normalized_text = self.old_text  # FIXME: Update this
+        self.utterance.oovs = " ".join(oovs)
+        self.utterance.ignored = not self.old_text
+        session.merge(self.utterance)
 
     def id(self) -> int:
         return 1
 
     def mergeWith(self, other: UpdateUtteranceTextCommand) -> bool:
-        if other.id() != self.id() or other.utterance_id != self.utterance_id:
+        if other.id() != self.id() or other.utterance.id != self.utterance.id:
             return False
         self.new_text = other.new_text
         return True
@@ -627,16 +597,17 @@ class UpdateSpeakerCommand(SpeakerCommand):
         )
 
 
-class UpdateUtteranceSpeakerCommand(CorpusCommand):
+class UpdateUtteranceSpeakerCommand(FileCommand):
     def __init__(
         self,
         utterances: typing.Union[Utterance, typing.List[Utterance]],
         new_speaker: typing.Union[Speaker, int],
-        corpus_model: CorpusModel,
+        file_model: FileUtterancesModel,
     ):
-        super().__init__(corpus_model)
+        super().__init__(file_model)
         if not isinstance(utterances, list):
             utterances = [utterances]
+        self.utterances = utterances
         self.utterance_ids = [x.id for x in utterances]
         self.file_ids = set([x.file_id for x in utterances])
         self.old_speaker_ids = [x.speaker_id for x in utterances]
@@ -644,7 +615,6 @@ class UpdateUtteranceSpeakerCommand(CorpusCommand):
             self.new_speaker_id = new_speaker
         else:
             self.new_speaker_id = new_speaker.id
-        self.resets_tier = True
         self.setText(
             QtCore.QCoreApplication.translate(
                 "UpdateUtteranceSpeakerCommand", "Update utterance speaker"
@@ -673,18 +643,15 @@ class UpdateUtteranceSpeakerCommand(CorpusCommand):
                 )
             )
             session.flush()
-        session.query(Utterance).filter(Utterance.id.in_(self.utterance_ids)).update(
-            {Utterance.speaker_id: self.new_speaker_id}
-        )
+        for u in self.utterances:
+            u.speaker_id = self.new_speaker_id
         session.query(Speaker).filter(
             Speaker.id.in_(self.old_speaker_ids + [self.new_speaker_id])
         ).update({Speaker.modified: True})
 
     def _undo(self, session) -> None:
-        update_mappings = []
-        for i, u_id in enumerate(self.utterance_ids):
-            update_mappings.append({"id": u_id, "speaker_id": self.old_speaker_ids[i]})
-        bulk_update(self.corpus_model.session, Utterance, update_mappings)
+        for i, u in enumerate(self.utterances):
+            u.speaker_id = self.old_speaker_ids[i]
         session.query(Speaker).filter(
             Speaker.id.in_(self.old_speaker_ids + [self.new_speaker_id])
         ).update({Speaker.modified: True})
@@ -692,6 +659,8 @@ class UpdateUtteranceSpeakerCommand(CorpusCommand):
     def update_data(self):
         super().update_data()
         self.corpus_model.set_file_modified(self.file_ids)
+        self.corpus_model.change_speaker_table_utterances(self.utterances)
+        self.file_model.change_speaker_table_utterances(self.utterances)
         self.corpus_model.changeCommandFired.emit()
         self.corpus_model.update_data()
         self.corpus_model.runFunction.emit(

anchor/widgets.py

@@ -25,6 +25,8 @@ from anchor.models import (
     CorpusSelectionModel,
     DiarizationModel,
     DictionaryTableModel,
+    FileSelectionModel,
+    FileUtterancesModel,
     OovModel,
     SpeakerModel,
     TextFilterQuery,
@@ -61,7 +63,6 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
         self.max_time = None
         self.start_load_time = None
         self.min_time = None
-        self.corpus_model = None
         self.selection_model = None
         self.timer = QtCore.QTimer(self)
         self.timer.setInterval(1)
@@ -78,7 +79,6 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
         self._audio_output.setDevice(self.devices.defaultAudioOutput())
         self.setAudioOutput(self._audio_output)
         self.playbackStateChanged.connect(self.reset_position)
-        self.mediaStatusChanged.connect(self.update_load)
         self.fade_in_anim = QtCore.QPropertyAnimation(self._audio_output, b"volume")
         self.fade_in_anim.setDuration(10)
         self.fade_in_anim.setStartValue(0.1)
@@ -95,11 +95,6 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
         self.fade_out_anim.finished.connect(super().pause)
         self.file_path = None
 
-    def update_load(self, state):
-        if state == self.MediaStatus.LoadedMedia:
-            self.reset_position()
-            self.audioReady.emit(True)
-
     def handle_error(self, *args):
         print("ERROR")
         print(args)
@@ -118,12 +113,22 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
         self.fade_in_anim.start()
 
     def startTime(self):
-        if self.selection_model.selected_min_time is not None:
+        if (
+            self.selection_model.selected_min_time is not None
+            and self.selection_model.min_time
+            <= self.selection_model.selected_min_time
+            <= self.selection_model.max_time
+        ):
             return self.selection_model.selected_min_time
         return self.selection_model.min_time
 
     def maxTime(self):
-        if self.selection_model.selected_max_time is not None:
+        if (
+            self.selection_model.selected_max_time is not None
+            and self.selection_model.min_time
+            <= self.selection_model.selected_max_time
+            <= self.selection_model.max_time
+        ):
             return self.selection_model.selected_max_time
         return self.selection_model.max_time
 
@@ -149,14 +154,10 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
                 break
         self._audio_output.setDevice(o)
 
-    def set_corpus_models(
-        self, corpus_model: Optional[CorpusModel], selection_model: Optional[CorpusSelectionModel]
-    ):
-        self.corpus_model = corpus_model
-        self.selection_model = selection_model
-        if corpus_model is None:
+    def set_models(self, selection_model: Optional[FileSelectionModel]):
+        if selection_model is None:
             return
-        # self.selection_model.fileAboutToChange.connect(self.unload_file)
+        self.selection_model = selection_model
         self.selection_model.fileChanged.connect(self.loadNewFile)
         self.selection_model.viewChanged.connect(self.update_times)
         self.selection_model.selectionAudioChanged.connect(self.update_selection_times)
@@ -187,29 +188,27 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
         self.setCurrentTime(self.startTime())
 
     def update_times(self):
-        if (
-            self.playbackState() == QtMultimedia.QMediaPlayer.PlaybackState.StoppedState
-            or self.currentTime() < self.startTime()
-            or self.currentTime() > self.maxTime()
-        ):
+        if self.currentTime() < self.startTime() or self.currentTime() > self.maxTime():
+            self.stop()
+        if self.playbackState() != QtMultimedia.QMediaPlayer.PlaybackState.PlayingState:
             self.setCurrentTime(self.startTime())
 
     def loadNewFile(self, *args):
         self.audioReady.emit(False)
         self.stop()
         try:
-            new_file = self.selection_model.current_file.sound_file.sound_file_path
+            new_file = self.selection_model.model().file.sound_file.sound_file_path
         except Exception:
             self.setSource(QtCore.QUrl())
             return
         if (
             self.selection_model.max_time is None
-            or self.selection_model.current_file is None
-            or self.selection_model.current_file.duration is None
+            or self.selection_model.model().file is None
+            or self.selection_model.model().file.duration is None
         ):
             self.setSource(QtCore.QUrl())
             return
-        self.channels = self.selection_model.current_file.num_channels
+        self.channels = self.selection_model.model().file.num_channels
         self.setSource(f"file:///{new_file}")
         self.setPosition(0)
         self.audioReady.emit(True)
@@ -218,19 +217,6 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
         pos = self.position()
         return pos / 1000
 
-    def setMaxTime(self, max_time):
-        if max_time is None:
-            return
-        self.max_time = max_time * 1000
-
-    def setMinTime(
-        self, min_time
-    ):  # Positions for MediaPlayer are in milliseconds, no SR required
-        if min_time is None:
-            min_time = 0
-        self.min_time = int(min_time * 1000)
-        self.setCurrentTime(min_time)
-
     def setCurrentTime(self, time):
         if time is None:
             time = 0
@@ -245,7 +231,7 @@ class MediaPlayer(QtMultimedia.QMediaPlayer):  # pragma: no cover
             self.stop()
             self.setSource(
                 QtCore.QUrl.fromLocalFile(
-                    self.selection_model.current_file.sound_file.sound_file_path
+                    self.selection_model.model().file.sound_file.sound_file_path
                 )
             )
             self.play()
@@ -416,7 +402,7 @@ class UtteranceListTable(AnchorTableView):
     def set_models(self, model: CorpusModel, selection_model: CorpusSelectionModel):
         self.setModel(model)
         self.setSelectionModel(selection_model)
-        self.doubleClicked.connect(self.selectionModel().focusUtterance)
+        self.doubleClicked.connect(self.selectionModel().focus_utterance)
         self.model().utteranceTextUpdated.connect(self.repaint)
         self.refresh_settings()
         model.corpusLoaded.connect(self.update_header)
@@ -793,6 +779,7 @@ class UtteranceDetailWidget(QtWidgets.QWidget):  # pragma: no cover
         self.settings = AnchorSettings()
         self.setAttribute(QtCore.Qt.WidgetAttribute.WA_StyledBackground, True)
         self.corpus_model = None
+        self.file_model = None
         self.selection_model = None
         self.dictionary_model = None
         self.plot_widget = UtteranceView(self)
@@ -829,24 +816,28 @@ class UtteranceDetailWidget(QtWidgets.QWidget):  # pragma: no cover
     def set_models(
         self,
         corpus_model: CorpusModel,
-        selection_model: CorpusSelectionModel,
+        file_model: FileUtterancesModel,
+        selection_model: FileSelectionModel,
         dictionary_model: DictionaryTableModel,
     ):
         self.corpus_model = corpus_model
+        self.file_model = file_model
         self.selection_model = selection_model
         self.dictionary_model = dictionary_model
         self.corpus_model.textFilterChanged.connect(self.plot_widget.set_search_term)
         self.selection_model.viewChanged.connect(self.update_to_slider)
         self.selection_model.fileChanged.connect(self.update_to_slider)
-        self.plot_widget.set_models(corpus_model, selection_model, self.dictionary_model)
+        self.plot_widget.set_models(
+            corpus_model, file_model, selection_model, self.dictionary_model
+        )
 
     def update_to_slider(self):
         with QtCore.QSignalBlocker(self.scroll_bar):
-            if self.selection_model.current_file is None or self.selection_model.min_time is None:
+            if self.selection_model.model().file is None or self.selection_model.min_time is None:
                 return
             if (
                 self.selection_model.min_time == 0
-                and self.selection_model.max_time == self.selection_model.current_file.duration
+                and self.selection_model.max_time == self.selection_model.model().file.duration
             ):
                 self.scroll_bar.setPageStep(10)
                 self.scroll_bar.setEnabled(False)
@@ -854,7 +845,7 @@ class UtteranceDetailWidget(QtWidgets.QWidget):  # pragma: no cover
                 self.pan_right_button.setEnabled(False)
                 self.scroll_bar.setMaximum(0)
                 return
-            duration_ms = int(self.selection_model.current_file.duration * 1000)
+            duration_ms = int(self.selection_model.model().file.duration * 1000)
             begin = self.selection_model.min_time * 1000
             end = self.selection_model.max_time * 1000
             window_size_ms = int(end - begin)