Skip to content

[DRAFT] Chunker .add API #2261

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: 05-29-anton_chunker
Choose a base branch
from
Draft

Conversation

atroyn
Copy link
Contributor

@atroyn atroyn commented May 30, 2024

Description of changes

This PR adds Chunkers to the .add method on Collection.

It allows users to pass a chunker in whenever the .add. The chunker splits documents into chunks according to supplied parameters.

This is implemented in a way that Chunkers can read off properties of collections - we want to do that so that Chunkers can be aware of embedding functions. (next PR)

Thinking about this a little more, we probably want max_tokens to already be a part of the Chunk interface, but I can implement that when I make these smarter in the next PR in this stack.

There is a notebook for you to try this out! Install chroma from the repo with

pip install -e .

Then check out https://app.graphite.dev/github/pr/chroma-core/chroma/2261/DRAFT-Chunker-add-API#file-chunker_example.ipynb

Test plan

TBC

Documentation Changes

TBC

TODO

  • Tests
  • Documentation

Copy link

vercel bot commented May 30, 2024

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name Status Preview Comments Updated (UTC)
chroma ✅ Ready (Inspect) Visit Preview 💬 Add feedback Jun 6, 2024 4:21am

Copy link

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

  • Can you think of any use case in which the code does not behave as intended? Have they been tested?
  • Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  • If appropriate, are there adequate property based tests?
  • If appropriate, are there adequate unit tests?
  • Should any logging, debugging, tracing information be added or removed?
  • Are error messages user-friendly?
  • Have all documentation changes needed been made?
  • Have all non-obvious changes been commented?

System Compatibility

  • Are there any potential impacts on other parts of the system or backward compatibility?
  • Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

  • Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

Copy link

Please tag your PR title with one of: [ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas

Copy link
Contributor Author

atroyn commented May 30, 2024

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @atroyn and the rest of your teammates on Graphite Graphite

@atroyn atroyn mentioned this pull request May 30, 2024
2 tasks
@atroyn atroyn changed the title Chunker .add API [DRAFT] Chunker .add API May 30, 2024
if chunker is not None:
chunks, chunk_ids = chunk_and_assign_ids(documents, ids, chunker)
# Flatten chunk_ids and chunks, and assign them as the new values
ids = [id for chunk_id in chunk_ids for id in chunk_id]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Basically we expand the list of IDs by chunk ID. We should probably pass these out to the user (we currently do not pass IDs out) so they can keep track of the new IDs. I can do that here, on in a separate PR.

@atroyn atroyn marked this pull request as ready for review May 30, 2024 04:09
embeddings = self._embed(input=documents)
elif images is not None:
if chunker is not None:
# Image chunking is not supported yet
raise NotImplementedError("Image chunking is not supported yet.")
embeddings = self._embed(input=images)
else:
if uris is None:
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Re: line 418]

`update` should be fairly straightforward since it's a delete and an add under the hood. We only need to do the same trick with IDs, delete the old ID and insert new ones.

See this comment inline on Graphite.

@HammadB
Copy link
Collaborator

HammadB commented May 30, 2024

I think my primary feedback is that:

  1. Adding chunk() to add() feels incongruous with the similar notions of DataLoaders and EmbeddingFunctions. Which are basically similar in the sense that they are pre-ingest transformations of the input to add(). For DataLoaders and EmbeddingFunctions, we make them a part of the collection.
  2. The issue with DataLoaders and EmbeddingFunctions is we don't persist them properly - which leads to weird bugs and user confusion since things like get_or_create can have the wrong one.

I think that we should either

  1. Fix the persistence issue and make the Chunker() part of the Collection. There are cheap ways for us to fix the persistence here IMO.
  2. Make the Chunker() not part of the API but just a function users can call and then call add() after. This sets us up to to be able to do more API level cleverness like making it stored with the collection or in a pipeline without adding complexity to the core API.
  3. Copy the same pattern as Efs/DLs

@atroyn atroyn force-pushed the 05-29-anton_chunker branch from d2b2d7f to 719c37c Compare May 31, 2024 19:56
@atroyn atroyn force-pushed the 05-29-chunker_.add_api branch from 19197fd to 14d8b1e Compare May 31, 2024 19:57
Copy link

Please tag your PR title with one of: [ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas

Copy link

github-actions bot commented Jun 3, 2024

Please tag your PR title with one of: [ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas

@atroyn atroyn force-pushed the 05-29-anton_chunker branch 2 times, most recently from d673865 to 205c06c Compare June 5, 2024 18:49
@atroyn atroyn force-pushed the 05-29-chunker_.add_api branch from aee51bf to bdd2cca Compare June 5, 2024 18:49
Copy link

github-actions bot commented Jun 5, 2024

Please tag your PR title with one of: [ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas

@atroyn atroyn force-pushed the 05-29-anton_chunker branch from 205c06c to 514cf4d Compare June 6, 2024 04:17
@atroyn atroyn force-pushed the 05-29-chunker_.add_api branch from bdd2cca to e9dbd2d Compare June 6, 2024 04:18
Copy link

github-actions bot commented Jun 6, 2024

Please tag your PR title with one of: [ENH | BUG | DOC | TST | BLD | PERF | TYP | CLN | CHORE]. See https://docs.trychroma.com/contributing#contributing-code-and-ideas

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants