Source code for finetuner.model
from dataclasses import dataclass
from typing import List, Union
from _finetuner.runner.stubs import model
from _finetuner.runner.stubs.model import * # noqa F401
from _finetuner.runner.stubs.model import _EmbeddingModelStub
[docs]def get_header() -> Tuple[str, ...]:
"""Get table header."""
return 'name', 'task', 'output_dim', 'architecture', 'description'
[docs]def get_row(model_stub) -> Tuple[str, ...]:
"""Get table row."""
return (
model_stub.display_name,
model_stub.task,
str(model_stub.output_shape[1]),
model_stub.architecture,
model_stub.description,
)
[docs]def list_model_classes() -> Dict[str, ModelStubType]:
rv = {}
members = inspect.getmembers(model, inspect.isclass)
parent_class = _EmbeddingModelStub
for name, stub in members:
if (
name != 'MLPStub'
and not name.startswith('_')
and type(stub) != type
and issubclass(stub, parent_class)
):
rv[name] = stub
return rv
[docs]@dataclass
class SynthesisModels:
"""Class specifying the models to be used in a data synthesis job.
:param: relation_miner: The name of the model or list of models to use for
relation mining.
:param cross_encoder: The name of the model to use as the cross encoder
"""
relation_miner: Union[str, List[str]]
cross_encoder: str
synthesis_model_en = SynthesisModels(
relation_miner='sbert-base-en',
cross_encoder='crossencoder-base-en',
)
synthesis_model_multi = SynthesisModels(
relation_miner='distiluse-base-multi',
cross_encoder='crossencoder-base-ml',
)