Skip to content

Commit

Permalink
chore(hf): remove unused arguments (#556)
Browse files Browse the repository at this point in the history
* chore(hf): remove unused arguments

* cleanup types a bit in hf
  • Loading branch information
shcheklein authored Nov 3, 2024
1 parent e1d7a25 commit 4f0713d
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 7 deletions.
2 changes: 1 addition & 1 deletion src/datachain/lib/dc.py
Original file line number Diff line number Diff line change
Expand Up @@ -1690,7 +1690,7 @@ def from_hf(

model_name = model_name or object_name or ""
hf_features = next(iter(ds_dict.values())).features
output = output | get_output_schema(hf_features, model_name)
output = output | get_output_schema(hf_features)
model = dict_to_data_model(model_name, output)
if object_name:
output = {object_name: model}
Expand Down
10 changes: 4 additions & 6 deletions src/datachain/lib/hf.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,17 +138,15 @@ def convert_feature(val: Any, feat: Any, anno: Any) -> Any: # noqa: PLR0911
return HFAudio(**val)


def get_output_schema(
features: Features, model_name: str = "", stream: bool = True
) -> dict[str, DataType]:
def get_output_schema(features: Features) -> dict[str, DataType]:
"""Generate UDF output schema from huggingface datasets features."""
fields_dict = {}
for name, val in features.items():
fields_dict[name] = _feature_to_chain_type(name, val) # type: ignore[assignment]
return fields_dict # type: ignore[return-value]
fields_dict[name] = _feature_to_chain_type(name, val)
return fields_dict


def _feature_to_chain_type(name: str, val: Any) -> type: # noqa: PLR0911
def _feature_to_chain_type(name: str, val: Any) -> DataType: # noqa: PLR0911
if isinstance(val, Value):
return arrow_type_mapper(val.pa_type)
if isinstance(val, ClassLabel):
Expand Down

0 comments on commit 4f0713d

Please sign in to comment.