phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-10 04:08:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix iterrows

Browse source
This commit is contained in:
Xi Yan 2025-03-15 14:24:46 -07:00
parent 82ec0d24f3
commit a197101635
2 changed files with 10 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

25
llama_stack/providers/inline/datasetio/localfs/datasetio.py
View file

@ -131,33 +131,24 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
async def iterrows( async def iterrows(
self, self,
dataset_id: str, dataset_id: str,
rows_in_page: int, start_index: Optional[int] = None,
page_token: Optional[str] = None, limit: Optional[int] = None,
filter_condition: Optional[str] = None,
) -> IterrowsResponse: ) -> IterrowsResponse:
dataset_info = self.dataset_infos.get(dataset_id) dataset_info = self.dataset_infos.get(dataset_id)
dataset_info.dataset_impl.load() dataset_info.dataset_impl.load()
if page_token and not page_token.isnumeric(): start_index = start_index or 0
raise ValueError("Invalid page_token")
if page_token is None or len(page_token) == 0: if limit is None or limit == -1:
next_page_token = 0
else:
next_page_token = int(page_token)
start = next_page_token
if rows_in_page == -1:
end = len(dataset_info.dataset_impl) end = len(dataset_info.dataset_impl)
else: else:
end = min(start + rows_in_page, len(dataset_info.dataset_impl)) end = min(start_index + limit, len(dataset_info.dataset_impl))
rows = dataset_info.dataset_impl[start:end] rows = dataset_info.dataset_impl[start_index:end]
return IterrowsResponse( return IterrowsResponse(
rows=rows, data=rows,
total_count=len(rows), next_index=end if end < len(dataset_info.dataset_impl) else None,
next_page_token=str(end),
) )
async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None: async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:

4
llama_stack/providers/remote/datasetio/huggingface/huggingface.py
View file

@ -84,7 +84,7 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
start_index = start_index or 0 start_index = start_index or 0
if limit == -1: if limit is None or limit == -1:
end = len(loaded_dataset) end = len(loaded_dataset)
else: else:
end = min(start_index + limit, len(loaded_dataset)) end = min(start_index + limit, len(loaded_dataset))
@ -93,7 +93,7 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
return IterrowsResponse( return IterrowsResponse(
data=rows, data=rows,
next_index=end, next_index=end if end < len(loaded_dataset) else None,
) )
async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None: async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None: