<commit_before>''' Extractor classes based on pre-trained models. '''

import numpy as np
import pandas as pd

from pliers.extractors.image import ImageExtractor
from pliers.extractors.base import Extractor, ExtractorResult
from pliers.filters.image import ImageResizingFilter
from pliers.stimuli import ImageStim, TextStim, AudioStim
from pliers.utils import (attempt_to_import, verify_dependencies,
                         listify)

import logging

tf = attempt_to_import('tensorflow')
hub = attempt_to_import('tensorflow_hub')
attempt_to_import('tensorflow.keras')
attempt_to_import('tensorflow_text')


class TFHubExtractor(Extractor):

    ''' A generic class for Tensorflow Hub extractors 
    Args:
        url_or_path (str): url or path to TFHub model. You can
            browse models at https://tfhub.dev/.
        task (str): model task/domain identifier
        features (optional): list of labels (for classification) 
            or other feature names. If a one-element list or 
            a string is passed, packs output vectors into one column.
            If not specified, returns numbered features 
            (feature_0, feature_1, ... ,feature_n).
        kwargs (dict): arguments to hub.KerasLayer call
    '''

    _log_attributes = ('url_or_path', 'features', 'transform_out')

    def __init__(self, url_or_path, features=None, task=None, 
                 transform_out=None, **kwargs):
        verify_dependencies(['tensorflow', 'tensorflow_hub', 
                             'tensorflow_text'])
        self.model = hub.KerasLayer(url_or_path, **kwargs)
        self.url_or_path = url_or_path
        self.features = listify(features)
        self._task = task
        self.transform_out = transform_out
        super().__init__()

    def get_feature_names(self, out):
        features = ['feature_' + str(i) 
                    for i in range(out.shape[-1])]
        return features
    
    def _preprocess(self, stim):
        if type(stim) == TextStim:
            return listify(stim.data)
        else:
            return stim.data

    def _postprocess(self, out):
        if self.transform_out:
            out = self.transform_out(out)
        return out.numpy().squeeze()<commit_msg>add Stim as input_type<commit_after>''' Extractor classes based on pre-trained models. '''

import numpy as np
import pandas as pd

from pliers.extractors.image import ImageExtractor
from pliers.extractors.base import Extractor, ExtractorResult
from pliers.filters.image import ImageResizingFilter
from pliers.stimuli import ImageStim, TextStim
from pliers.stimuli.base import Stim
from pliers.utils import (attempt_to_import, verify_dependencies,
                         listify)

import logging

tf = attempt_to_import('tensorflow')
hub = attempt_to_import('tensorflow_hub')
attempt_to_import('tensorflow.keras')
attempt_to_import('tensorflow_text')


class TFHubExtractor(Extractor):

    ''' A generic class for Tensorflow Hub extractors 
    Args:
        url_or_path (str): url or path to TFHub model. You can
            browse models at https://tfhub.dev/.
        task (str): model task/domain identifier
        features (optional): list of labels (for classification) 
            or other feature names. If a one-element list or 
            a string is passed, packs output vectors into one column.
            If not specified, returns numbered features 
            (feature_0, feature_1, ... ,feature_n).
        transform_out (optional): function to transform model 
            output for compatibility with extractor result
        transform_inp (optional): function to transform Stim.data 
            for compatibility with model input format
        kwargs (dict): arguments to hub.KerasLayer call
    '''

    _log_attributes = ('url_or_path', 'features', 'transform_out')
    _input_type = Stim

    def __init__(self, url_or_path, features=None,
                 transform_out=None, transform_inp=None,
                  **kwargs):
        verify_dependencies(['tensorflow', 'tensorflow_hub', 
                             'tensorflow_text'])
        self.model = hub.KerasLayer(url_or_path, **kwargs)
        self.url_or_path = url_or_path
        self.features = listify(features)
        self.transform_out = transform_out
        self.transform_inp = transform_inp
        super().__init__()

    def get_feature_names(self, out):
        features = ['feature_' + str(i) 
                    for i in range(out.shape[-1])]
        return features
    
    def _preprocess(self, stim):
        if self.transform_inp:
            return self.transform_inp(stim.data)
        else:
            if type(stim) == TextStim:
                return listify(stim.data)
            else:
                return stim.data

    def _postprocess(self, out):
        if self.transform_out:
            out = self.transform_out(out)
        return out.numpy().squeeze()
{'fraction_non_alphanumeric': 0.08424908424908426, 'fraction_numerical': 0.001221001221001221, 'mean_word_length': 2.5823615160349855, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '93035', 'n_tokens_mistral': 1396, 'n_tokens_neox': 1327, 'n_words': 388}