Improve multi modul support and refactor

This commit is contained in:
mesonium
2024-07-02 20:05:45 +02:00
committed by hueso
parent aff69e5b40
commit b0adecef7a
6 changed files with 111 additions and 92 deletions
+90 -69
View File
@@ -15,6 +15,7 @@
from __future__ import annotations
from dataclasses import asdict
import logging
import typing
from pathlib import Path
@@ -40,45 +41,51 @@ if TYPE_CHECKING:
log = logging.getLogger('gajim.p.sttvm_config_dialog')
SUPPORTED_MODELS: dict[str, dict[str, typing.Union[list[str], Any, str]]] = {
'model_openaiwhisper': {
'moduls': ['whisper'],
'class': openai_whisper.WhisperModel,
'name': 'OpenAI Whisper'
},
'model_ctranslate2': {
'moduls': ['ctranslate2'],
'class': None,
'name': _('CTranslate2')
},
'model_faster-whisper': {
'moduls': ['faster-whisper'],
'class': None,
'name:': _('Faster-Whisper')
},
'model_distill': {
'moduls': ['transformers', 'accelerate', 'datasets[audio]'],
'class': None,
'name': _('Distill')
}
@dataclass
class Model:
name: str
required_moduls: list[str]
klass: object
config: Any
instance: typing.Optional[object] = None
SUPPORTED_MODELS: dict[str, Model] = {
'model_openaiwhisper': Model('OpenAI Whisper',
['whisper'],
openai_whisper.WhisperModel,
OpenAIWhisperSettings),
'model_ctranslate2': Model('CTranslate2',
['ctranslate2'],
None,
None),
'model_faster-whisper': Model('Fast-Whisper',
['faster-whisper'],
None,
None),
'model_distill': Model('Distill',
['transformers', 'accelerate', 'datasets[audio]'],
None,
None)
}
class Configuration:
def __init__(self, plugin: STTVoiceMessagesPlugin):
self._plugin = plugin
self._openaiwhisper_settings = OpenAIWhisperSettings()
self._available_models: dict[
str, dict[str, typing.Union[list[str], Any, str]]] = {}
self._available_models: dict[str, Model] = {}
self.check_available_moduls()
log.debug('config = %s', self._plugin.config['model_openaiwhisper'])
@property
def plugin(self) -> STTVoiceMessagesPlugin:
return self._plugin
@property
def available_models(self) -> dict[
str, dict[str, typing.Union[list[str], Any, str]]]:
def available_models(self) -> dict[str, Model]:
return self._available_models
def on_setting(self, value: Any, data: Any) -> None:
@@ -86,58 +93,61 @@ class Configuration:
value.strip()
log.debug('plugin config before:\n %s', self.plugin.config.data)
# TODO: Is 'modelname_key = value' a good design?
self.plugin.config[data] = value
# TODO: Apply setting only to specific instance
self._plugin.config['model_instance'].on_setting(data, value)
log.debug('plugin config after:\n %s', self.plugin.config.data)
def on_set_model(self, value: Any, data: Any) -> None:
def on_config_model(self, model: str, value: Any, data: Any) -> None:
if isinstance(value, str):
value.strip()
log.debug('plugin config before:\n %s', self.plugin.config.data[model])
setattr(self.plugin.config.data[model], data, value)
log.debug('plugin config after:\n %s', self.plugin.config.data[model])
self._plugin.config.data[model].instance.set_config(self.plugin.config.data[model])
def on_set_model(self, model: Any) -> None:
if isinstance(model, str):
model.strip()
log.debug('plugin config before:\n %s', self.plugin.config.data)
self._available_models[value]['model_instance'] = \
self._available_models[value]['class']()
if (self.plugin.config.data[model].instance is None and
self._available_models[model].klass is not None):
self.plugin.config.data[model].instance = \
self._available_models[model].klass()
else:
return
self.plugin.config['model_class'] = self._available_models[value][
'class']
self.plugin.config['model_instance'] = self._available_models[value][
'model_instance']
self.on_setting(value, data)
self.plugin.config['model'] = model
log.debug('plugin config after:\n %s', self.plugin.config.data)
@staticmethod
def is_module_available(module: str) -> bool:
try:
__import__(module)
return True
except ModuleNotFoundError:
log.debug('Could not find module %s', module)
return False
except ImportError as ex:
log.debug(str(ex))
return False
def check_available_moduls(self):
def is_module_available(module: str) -> bool:
try:
__import__(module)
return True
except ModuleNotFoundError:
log.debug('Could not find module %s', module)
return False
except ImportError as ex:
log.debug(str(ex))
return False
for model in SUPPORTED_MODELS:
available = True
for modul in SUPPORTED_MODELS[model]['moduls']:
if not self.is_module_available(modul):
for modul in SUPPORTED_MODELS[model].required_moduls:
if not is_module_available(modul):
available = False
continue
if available:
self._available_models[model] = SUPPORTED_MODELS[model]
if SUPPORTED_MODELS[model].config is not None:
log.debug('created config for model = %s: %s', model, self._available_models[model])
log.debug('plugin config for model = %s', self.plugin.config[model])
self.plugin.config.data[model].instance = None
self._available_models[model].config = self.plugin.config[model]
if (self.plugin.config.data['model_class'] is None
and len(self._available_models) > 0):
model = list(self._available_models)[0]
self.on_set_model(model, 'model')
log.debug('Choose first available model!')
else:
log.debug('Available model already chosen!')
self.on_set_model(self._plugin.config['model'])
log.debug('models = %s', self._available_models)
@@ -188,6 +198,9 @@ class STTVoiceMessagesConfigDialog(Gtk.ApplicationWindow):
self.show_all()
############################################################################
# General Settings
############################################################################
class STTBehaviour(PreferenceBox):
def __init__(self, config_dialog: STTVoiceMessagesConfigDialog) -> None:
settings = [
@@ -206,7 +219,7 @@ class STTVoiceMessagesConfigDialog(Gtk.ApplicationWindow):
models: list[tuple[str, str]] = []
for key, value in config_dialog.config.available_models.items():
models.append(
(key, str(value['name']))
(key, str(value.name))
)
settings = [
@@ -222,29 +235,37 @@ class STTVoiceMessagesConfigDialog(Gtk.ApplicationWindow):
PreferenceBox.__init__(self, settings)
############################################################################
# OpenAI Whisper Settings
############################################################################
class OpenAIWhisperGeneral(PreferenceBox):
def __init__(self, config_dialog: STTVoiceMessagesConfigDialog) -> None:
self._model = 'model_openaiwhisper'
self._config_dialog = config_dialog
settings = [
Setting(SettingKind.POPOVER,
_('Language Model Size'),
SettingType.VALUE,
value=config_dialog.plugin.config[
'whisperai_model_size'],
data='whisperai_model_size',
callback=config_dialog.config.on_setting,
value=config_dialog.config.available_models[self._model].config.model_size,
data='model_size',
callback=self._set_config,
props={'entries': whisper.available_models()}),
Setting(SettingKind.SWITCH,
_('Translate'),
SettingType.VALUE,
value=config_dialog.plugin.config[
'whisperai_translate'],
data='whisperai_translate',
callback=config_dialog.config.on_setting)
value=config_dialog.config.available_models[self._model].config.translate_to_english,
data='translate_to_english',
callback=self._set_config)
]
PreferenceBox.__init__(self, settings)
def _set_config(self, value: Any, data: Any):
self._config_dialog.config.on_config_model(self._model, value, data)
def _add_prefs(self, prefs: list[tuple[str, type[PreferenceBox]]]):
for ui_name, klass in prefs:
pref_box = getattr(self._ui, ui_name)
+2 -1
View File
@@ -57,7 +57,8 @@ class STTBox(Gtk.Box):
def _on_transcribe_clicked(self, _button: Gtk.Button) -> None:
log.debug('config.data = %s', self._config.data)
model = self._config.data['model_instance']
model_name = self._config.data['model']
model = self._config.data[model_name].instance
if model is None:
return