mseep-txtai 9.1.1__py3-none-any.whl

txtai/models/registry.py

@@ -0,0 +1,37 @@
+"""
+Registry module
+"""
+
+from transformers import AutoModel, AutoModelForQuestionAnswering, AutoModelForSequenceClassification
+from transformers.models.auto.tokenization_auto import TOKENIZER_MAPPING
+
+
+class Registry:
+    """
+    Methods to register models and fully support pipelines.
+    """
+
+    @staticmethod
+    def register(model, config=None):
+        """
+        Registers a model with auto model and tokenizer configuration to fully support pipelines.
+
+        Args:
+            model: model to register
+            config: config class name
+        """
+
+        # Default config class to model class if not provided
+        config = config if config else model.__class__
+
+        # Default model config_class if empty
+        if hasattr(model.__class__, "config_class") and not model.__class__.config_class:
+            model.__class__.config_class = config
+
+        # Add references for this class to supported AutoModel classes
+        for mapping in [AutoModel, AutoModelForQuestionAnswering, AutoModelForSequenceClassification]:
+            mapping.register(config, model.__class__)
+
+        # Add references for this class to support pipeline AutoTokenizers
+        if hasattr(model, "config") and type(model.config) not in TOKENIZER_MAPPING:
+            TOKENIZER_MAPPING.register(type(model.config), type(model.config).__name__)

txtai/models/tokendetection.py

@@ -0,0 +1,122 @@
+"""
+Token Detection module
+"""
+
+import inspect
+import os
+
+import torch
+
+from transformers import PreTrainedModel
+
+
+class TokenDetection(PreTrainedModel):
+    """
+    Runs the replaced token detection training objective. This method was first proposed by the ELECTRA model.
+    The method consists of a masked language model generator feeding data to a discriminator that determines
+    which of the tokens are incorrect. More on this training objective can be found in the ELECTRA paper.
+    """
+
+    def __init__(self, generator, discriminator, tokenizer, weight=50.0):
+        """
+        Creates a new TokenDetection class.
+
+        Args:
+            generator: Generator model, must be a masked language model
+            discriminator: Discriminator model, must be a model that can detect replaced tokens. Any model can
+                           can be customized for this task. See ElectraForPretraining for more.
+        """
+
+        # Initialize model with discriminator config
+        super().__init__(discriminator.config)
+
+        self.generator = generator
+        self.discriminator = discriminator
+
+        # Tokenizer to save with generator and discriminator
+        self.tokenizer = tokenizer
+
+        # Discriminator weight
+        self.weight = weight
+
+        # Share embeddings if both models are the same type
+        # Embeddings must be same size
+        if self.generator.config.model_type == self.discriminator.config.model_type:
+            self.discriminator.set_input_embeddings(self.generator.get_input_embeddings())
+
+        # Set attention mask present flags
+        self.gattention = "attention_mask" in inspect.signature(self.generator.forward).parameters
+        self.dattention = "attention_mask" in inspect.signature(self.discriminator.forward).parameters
+
+    # pylint: disable=E1101
+    def forward(self, input_ids=None, labels=None, attention_mask=None, token_type_ids=None):
+        """
+        Runs a forward pass through the model. This method runs the masked language model then randomly samples
+        the generated tokens and builds a binary classification problem for the discriminator (detecting if each token is correct).
+
+        Args:
+            input_ids: token ids
+            labels: token labels
+            attention_mask: attention mask
+            token_type_ids: segment token indices
+
+        Returns:
+            (loss, generator outputs, discriminator outputs, discriminator labels)
+        """
+
+        # Copy input ids
+        dinputs = input_ids.clone()
+
+        # Run inputs through masked language model
+        inputs = {"attention_mask": attention_mask} if self.gattention else {}
+        goutputs = self.generator(input_ids, labels=labels, token_type_ids=token_type_ids, **inputs)
+
+        # Get predictions
+        preds = torch.softmax(goutputs[1], dim=-1)
+        preds = preds.view(-1, self.config.vocab_size)
+
+        tokens = torch.multinomial(preds, 1).view(-1)
+        tokens = tokens.view(dinputs.shape[0], -1)
+
+        # Labels have a -100 value to ignore loss from unchanged tokens
+        mask = labels.ne(-100)
+
+        # Replace the masked out tokens of the input with the generator predictions
+        dinputs[mask] = tokens[mask]
+
+        # Turn mask into new target labels - 1 (True) for corrupted, 0 otherwise.
+        # If the prediction was correct, mark it as uncorrupted.
+        correct = tokens == labels
+        dlabels = mask.long()
+        dlabels[correct] = 0
+
+        # Run token classification, predict whether each token was corrupted
+        inputs = {"attention_mask": attention_mask} if self.dattention else {}
+        doutputs = self.discriminator(dinputs, labels=dlabels, token_type_ids=token_type_ids, **inputs)
+
+        # Compute combined loss
+        loss = goutputs[0] + self.weight * doutputs[0]
+        return loss, goutputs[1], doutputs[1], dlabels
+
+    def save_pretrained(self, output, state_dict=None, **kwargs):
+        """
+        Saves current model to output directory.
+
+        Args:
+            output: output directory
+            state_dict: model state
+            kwargs: additional keyword arguments
+        """
+
+        # Save combined model to support training from checkpoints
+        super().save_pretrained(output, state_dict, **kwargs)
+
+        # Save generator tokenizer and model
+        gpath = os.path.join(output, "generator")
+        self.tokenizer.save_pretrained(gpath)
+        self.generator.save_pretrained(gpath)
+
+        # Save discriminator tokenizer and model
+        dpath = os.path.join(output, "discriminator")
+        self.tokenizer.save_pretrained(dpath)
+        self.discriminator.save_pretrained(dpath)

txtai/pipeline/__init__.py

@@ -0,0 +1,17 @@
+"""
+Pipeline imports
+"""
+
+from .audio import *
+from .base import Pipeline
+from .data import *
+from .factory import PipelineFactory
+from .hfmodel import HFModel
+from .hfpipeline import HFPipeline
+from .image import *
+from .llm import *
+from .llm import RAG as Extractor
+from .nop import Nop
+from .text import *
+from .tensors import Tensors
+from .train import *

txtai/pipeline/audio/__init__.py

@@ -0,0 +1,11 @@
+"""
+Audio imports
+"""
+
+from .audiomixer import AudioMixer
+from .audiostream import AudioStream
+from .microphone import Microphone
+from .signal import Signal
+from .texttoaudio import TextToAudio
+from .texttospeech import TextToSpeech
+from .transcription import Transcription

txtai/pipeline/audio/audiomixer.py

@@ -0,0 +1,58 @@
+"""
+AudioMixer module
+"""
+
+from ..base import Pipeline
+from .signal import Signal, SCIPY
+
+
+class AudioMixer(Pipeline):
+    """
+    Mixes multiple audio streams into a single stream.
+    """
+
+    def __init__(self, rate=None):
+        """
+        Creates an AudioMixer pipeline.
+
+        Args:
+            rate: optional target sample rate, otherwise uses input target rate with each audio segment
+        """
+
+        if not SCIPY:
+            raise ImportError('AudioMixer pipeline is not available - install "pipeline" extra to enable.')
+
+        # Target sample rate
+        self.rate = rate
+
+    def __call__(self, segment, scale1=1, scale2=1):
+        """
+        Mixes multiple audio streams into a single stream.
+
+        Args:
+            segment: ((audio1, sample rate), (audio2, sample rate))|list
+            scale1: optional scaling factor for segment1
+            scale2: optional scaling factor for segment2
+
+        Returns:
+            list of (audio, sample rate)
+        """
+
+        # Convert single element to list
+        segments = [segment] if isinstance(segment, tuple) else segment
+
+        results = []
+        for segment1, segment2 in segments:
+            audio1, rate1 = segment1
+            audio2, rate2 = segment2
+
+            # Resample audio, as necessary
+            target = self.rate if self.rate else rate1
+            audio1 = Signal.resample(audio1, rate1, target)
+            audio2 = Signal.resample(audio2, rate2, target)
+
+            # Mix audio into single segment
+            results.append((Signal.mix(audio1, audio2, scale1, scale2), target))
+
+        # Return single element if single element passed in
+        return results[0] if isinstance(segment, tuple) else results

txtai/pipeline/audio/audiostream.py

@@ -0,0 +1,94 @@
+"""
+AudioStream module
+"""
+
+from queue import Queue
+from threading import Thread
+
+# Conditional import
+try:
+    import sounddevice as sd
+
+    from .signal import Signal, SCIPY
+
+    AUDIOSTREAM = SCIPY
+except (ImportError, OSError):
+    AUDIOSTREAM = False
+
+from ..base import Pipeline
+
+
+class AudioStream(Pipeline):
+    """
+    Threaded pipeline that streams audio segments to an output audio device. This pipeline is designed
+    to run on local machines given that it requires access to write to an output device.
+    """
+
+    # End of stream message
+    COMPLETE = (1, None)
+
+    def __init__(self, rate=None):
+        """
+        Creates an AudioStream pipeline.
+
+        Args:
+            rate: optional target sample rate, otherwise uses input target rate with each audio segment
+        """
+
+        if not AUDIOSTREAM:
+            raise ImportError(
+                (
+                    'AudioStream pipeline is not available - install "pipeline" extra to enable. '
+                    "Also check that the portaudio system library is available."
+                )
+            )
+
+        # Target sample rate
+        self.rate = rate
+
+        self.queue = Queue()
+        self.thread = Thread(target=self.play)
+        self.thread.start()
+
+    def __call__(self, segment):
+        """
+        Queues audio segments for the audio player.
+
+        Args:
+            segment: (audio, sample rate)|list
+
+        Returns:
+            segment
+        """
+
+        # Convert single element to list
+        segments = [segment] if isinstance(segment, tuple) else segment
+
+        for x in segments:
+            self.queue.put(x)
+
+        # Return single element if single element passed in
+        return segments[0] if isinstance(segment, tuple) else segments
+
+    def wait(self):
+        """
+        Waits for all input audio segments to be played.
+        """
+
+        self.thread.join()
+
+    def play(self):
+        """
+        Reads audio segments from queue. This method runs in a separate non-blocking thread.
+        """
+
+        audio, rate = self.queue.get()
+        while not isinstance(audio, int) or (audio, rate) != AudioStream.COMPLETE:
+            # Resample to target sample rate, if necessary
+            audio, rate = (Signal.resample(audio, rate, self.rate), self.rate) if self.rate else (audio, rate)
+
+            # Play audio segment
+            sd.play(audio, rate, blocking=True)
+
+            # Get next segment
+            audio, rate = self.queue.get()

txtai/pipeline/audio/microphone.py

@@ -0,0 +1,244 @@
+"""
+Microphone module
+"""
+
+import logging
+
+import numpy as np
+
+# Conditional import
+try:
+    import sounddevice as sd
+    import webrtcvad
+
+    from scipy.signal import butter, sosfilt
+
+    from .signal import Signal, SCIPY
+
+    MICROPHONE = SCIPY
+except (ImportError, OSError):
+    MICROPHONE = False
+
+from ..base import Pipeline
+
+# Logging configuration
+logger = logging.getLogger(__name__)
+
+
+class Microphone(Pipeline):
+    """
+    Reads input speech from a microphone device. This pipeline is designed to run on local machines given
+    that it requires access to read from an input device.
+    """
+
+    def __init__(self, rate=16000, vadmode=3, vadframe=20, vadthreshold=0.6, voicestart=300, voiceend=3400, active=5, pause=8):
+        """
+        Creates a new Microphone pipeline.
+
+        Args:
+            rate: sample rate to record audio in, defaults to 16000 (16 kHz)
+            vadmode: aggressiveness of the voice activity detector (1 - 3), defaults to 3, which is the most aggressive filter
+            vadframe: voice activity detector frame size in ms, defaults to 20
+            vadthreshold: percentage of frames (0.0 - 1.0) that must be voice to be considered speech, defaults to 0.6
+            voicestart: starting frequency to use for voice filtering, defaults to 300
+            voiceend: ending frequency to use for voice filtering, defaults to 3400
+            active: minimum number of active speech chunks to require before considering this speech, defaults to 5
+            pause: number of non-speech chunks to keep before considering speech complete, defaults to 8
+        """
+
+        if not MICROPHONE:
+            raise ImportError(
+                (
+                    'Microphone pipeline is not available - install "pipeline" extra to enable. '
+                    "Also check that the portaudio system library is available."
+                )
+            )
+
+        # Sample rate
+        self.rate = rate
+
+        # Voice activity detector
+        self.vad = webrtcvad.Vad(vadmode)
+        self.vadframe = vadframe
+        self.vadthreshold = vadthreshold
+
+        # Voice spectrum
+        self.voicestart = voicestart
+        self.voiceend = voiceend
+
+        # Audio chunks counts
+        self.active = active
+        self.pause = pause
+
+    def __call__(self, device=None):
+        """
+        Reads audio from an input device.
+
+        Args:
+            device: optional input device id, otherwise uses system default
+
+        Returns:
+            list of (audio, sample rate)
+        """
+
+        # Listen for audio
+        audio = self.listen(device[0] if isinstance(device, list) else device)
+
+        # Return single element if single element passed in
+        return (audio, self.rate) if device is None or not isinstance(device, list) else [(audio, self.rate)]
+
+    def listen(self, device):
+        """
+        Listens for speech. Detected speech is converted to 32-bit floats for compatibility with
+        automatic speech recognition (ASR) pipelines.
+
+        This method blocks until speech is detected.
+
+        Args:
+            device: input device
+
+        Returns:
+            audio
+        """
+
+        # Record in 100ms chunks
+        chunksize = self.rate // 10
+
+        # Open input stream
+        stream = sd.RawInputStream(device=device, samplerate=self.rate, channels=1, blocksize=chunksize, dtype=np.int16)
+
+        # Start the input stream
+        stream.start()
+
+        record, speech, nospeech, chunks = True, 0, 0, []
+        while record:
+            # Read chunk
+            chunk, _ = stream.read(chunksize)
+
+            # Detect speech using WebRTC VAD for audio chunk
+            detect = self.detect(chunk)
+            speech = speech + 1 if detect else speech
+            nospeech = 0 if detect else nospeech + 1
+
+            # Save chunk, if this is an active stream
+            if speech:
+                chunks.append(chunk)
+
+                # Pause limit has been reached, check if this audio should be accepted
+                if nospeech >= self.pause:
+                    logger.debug("Audio detected and being analyzed")
+                    if speech >= self.active and self.isspeech(chunks[:-nospeech]):
+                        # Disable recording
+                        record = False
+                    else:
+                        # Reset parameters and keep recording
+                        logger.debug("Speech not detected")
+                        speech, nospeech, chunks = 0, 0, []
+
+        # Stop the input stream
+        stream.stop()
+
+        # Convert to float32 and return
+        audio = np.frombuffer(b"".join(chunks), np.int16)
+        return Signal.float32(audio)
+
+    def isspeech(self, chunks):
+        """
+        Runs an ensemble of Voice Activity Detection (VAD) methods. Returns true if speech is
+        detected in the input audio chunks.
+
+        Args:
+            chunks: input audio chunks as byte buffers
+
+        Returns:
+            True if speech is detected, False otherwise
+        """
+
+        # Convert to NumPy array for processing
+        audio = np.frombuffer(b"".join(chunks), dtype=np.int16)
+
+        # Ensemble of:
+        #  - WebRTC VAD with a human voice range butterworth bandpass filter applied to the signal
+        #  - FFT applied to detect the energy ratio for human voice range vs total range
+        return self.detectband(audio) and self.detectenergy(audio)
+
+    def detect(self, buffer):
+        """
+        Detect speech using the WebRTC Voice Activity Detector (VAD).
+
+        Args:
+            buffer: input audio buffer frame as bytes
+
+        Returns:
+            True if the number of audio frames with audio pass vadthreshold, False otherwise
+        """
+
+        n = int(self.rate * (self.vadframe / 1000.0) * 2)
+        offset = 0
+
+        detects = []
+        while offset + n <= len(buffer):
+            detects.append(1 if self.vad.is_speech(buffer[offset : offset + n], self.rate) else 0)
+            offset += n
+
+        # Calculate detection ratio and return
+        ratio = sum(detects) / len(detects) if detects else 0
+        if ratio > 0:
+            logger.debug("DETECT %.4f", ratio)
+
+        return ratio >= self.vadthreshold
+
+    def detectband(self, audio):
+        """
+        Detects speech using audio data filtered through a butterworth band filter
+        with the human voice range.
+
+        Args:
+            audio: input audio data as an NumPy array
+
+        Returns:
+            True if speech is detected, False otherwise
+        """
+
+        # Upsample to float32
+        audio = Signal.float32(audio)
+
+        # Human voice frequency range
+        low = self.voicestart / (0.5 * self.rate)
+        high = self.voiceend / (0.5 * self.rate)
+
+        # Low and high pass filter using human voice range
+        sos = butter(5, Wn=[low, high], btype="band", output="sos")
+        audio = sosfilt(sos, audio)
+
+        # Scale back to int16
+        audio = Signal.int16(audio)
+
+        # Pass filtered signal to WebRTC VAD
+        return self.detect(audio.tobytes())
+
+    def detectenergy(self, audio):
+        """
+        Detects speech by comparing the signal energy of the human voice range
+        to the overall signal energy.
+
+        Args:
+            audio: input audio data as an NumPy array
+
+        Returns:
+            True if speech is detected, False otherwise
+        """
+
+        # Calculate signal energy
+        energyfreq = Signal.energy(audio, self.rate)
+
+        # Sum speech energy
+        speechenergy = 0
+        for f, e in energyfreq.items():
+            if self.voicestart <= f <= self.voiceend:
+                speechenergy += e
+
+        # Calculate ratio of speech energy to total energy and return
+        ratio = speechenergy / sum(energyfreq.values())
+        logger.debug("SPEECH %.4f", ratio)
+        return ratio >= self.vadthreshold