Suggested architecture for Kedro & LLMs

[astrojuanlu]

I've been toying with GenAI & Kedro for a while, first with Hugging Face (thanks to our datasets) and now through Ollama, which gives an OpenAI-compatible HTTP API:

$ curl http://localhost:11434/api/generate -d '{
  "model": "llama3",
  "prompt":"Why is the sky blue?"
}'

I'm trying to wrap my head around what would be the "blessed" way to implement this in Kedro.

1. Frozen inputs

One way would be to let the dataset include the prompt as part of the configuration, as follows:

llama3_ds:
  type: OllamaDataset
  model: llama3
  prompt: "Why is the sky blue?"

so catalog.load("llama3_ds") gives immediately an output.

Pro: The Kedro dataset performs all of the I/O needed. It's probably the "most pure" implementation. Could be modeled by using the already existing APIDataset.
Con: What if my prompt is not known at definition time?

2. Wrapper object

Another way of doing it would be returning a wrapper object to the node:

llama3_ds:
  type: OllamaDataset
  model: llama3

Which would then allow me to do things like

node(
  func=get_response,
  inputs=["llama3_ds", "param:prompt1"],
)

...

def get_response(llm, prompt: str):
  return llm.call(prompt)

Pro: It allows dynamic prompts, which might be necessary for some use cases.
Con 1: The node is performing I/O, which means that we're "breaking" the Kedro notion of what the node and the dataset should do.
Con 2: It breaks reproducibility, given that the prompt is not static (in this case it was a param, but could as well be the output of another node)

At the same time, we're already "breaking" that I/O separation in several places:

• PartitionedDataset https://docs.kedro.org/en/stable/data/partitioned_and_incremental_datasets.html#partitioned-dataset-load
• Our LangChain datasets https://docs.kedro.org/projects/kedro-datasets/en/latest/api/kedro_datasets_experimental.langchain.ChatOpenAIDataset.html#kedro_datasets_experimental.langchain.ChatOpenAIDataset

Any thoughts?

[lordsoffallen]

Hi @astrojuanlu,

Coming from our PR to share my thoughts on how kedro should work with LLMs.

My usual workflow is LLM is like this: I need a open source model so that requires defining some params for huggingface model. With LLMs in mind, there is lots of configs around managing models which I believe are suited to kedro catalog but I don't personally like considering them as datasets. So some separation like data catalog and model catalog would be highly appreciated. I am also not a fan of most of the claimed to be LLM tools our there, that includes langchain, llamaindex etc. They really aim for high level users with little to no experience in building ML models. I consider kedro as a tool for data scientists hence I believe they should be avoided early on. Could be optional dependency for certain cases but I don't see the need.

For huggingface I can imagine wrappers around getting model and tokenizer or pipeline depending on the use case is a good start. Also ability to save to local disk and/or to huggingface. This could be a dataset or a model as well. For instance, creating embeddings from a dataset and I would like to save them to huggingface. Adding more vector store capabilities would be nice abstraction like both open source tools as well as some apis. (qdrant etc)

Now for the API endpoint, there were some nice packages ( i forgot the name) which makes a unified entrypoint for the all other different apis like OpenAI, Claude etc. This would def help however a this makes a call to an endpoint so here, useful features such as adding pricing or logging. Ability to test this with a smaller local model would be super nice (not sure how for now).

Lastly, there is some work around prompts. Managing prompts, strings etc. Maybe there is something to do there.

Overall since kedro aims to manage data IO, simplifying getting the open source models (even lazy loading big models), making it seamless to work with APIs is a step in the right direction imo.

Let me know if there are some points that weren't clear. Happy to explain it in detail.

[astrojuanlu]

Thanks a lot @lordsoffallen for adding your thoughts. There's a lot to unpack so I'll go piece by piece, just picking one for now:

With LLMs in mind, there is lots of configs around managing models which I believe are suited to kedro catalog but I don't personally like considering them as datasets.

Now that you mention it, it was one of the elements surfaced in @AlpAribal & @datajoely research on MLOps platforms https://github.com/kedro-org/kedro/wiki/Synthesis-of-research-related-to-deployment-of-Kedro-to-modern-MLOps-platforms#differentiating-between-data-model-and-reporting-artifacts

In the end we call them datasets for historical reasons, even though many of them aren't really datasets.

[lordsoffallen]

I mean for the standard data science workflow there weren't a need to do much as they were just pickle dataset in the end. So you can do regular imports and use whatever model you like and gets serialized in the same way. LLMs on the hand, can be just bunch of configs. Example of what I mean here:

Standard Flow: -> Works for linear regression, XGboost, LightGBM etc.

model:
    type: PickleDataset
    filepath: ...

LLM Flow

llama3:
    type: HFModel
    checkpoint: meta/llama-3

gemma-2:
   type: HFModel
   checkpoint: google/gemma2

As you can see if I wanna put down couple of models and eval them and pick one, multiple configs suffice but mixing with training data grows quickly. I could convert this into a parameter and try to use a unifed function that gets the right model but that also goes out of the kedro principle of doing data IO but me doing that manually instead. Hope the example extends on my point a bit more.

[datajoely]

I would love to make this metadata a first class citizen:

cars:
  type: pandas.CSVDataset
  filepath: data/01_raw/company/cars.csv
  load_args:
    sep: ','
  save_args:
    index: False
    date_format: '%Y-%m-%d %H:%M'
    decimal: .
   metadata:
     viz:
      layer: raw
+     kind: data # or (model/artifact)

We could take inspiration from Kubeflow's artifact types

[astrojuanlu]

I wanna put down couple of models and eval them and pick one

Indeed, Kedro is a bit too obtuse for parallel experimentation, there's a lot of context about that in #1606

https://getindata.com/blog/kedro-dynamic-pipelines/ discusses one possible approach for tackling that, we know lots of users get there but we don't have enough context on whether it's enough or not.

I could convert this into a parameter and try to use a unifed function that gets the right model but that also goes out of the kedro principle of doing data IO

👀 indeed, parametrising the catalog itself in this way is not always easy.

You could have a look at runtime parameters like these: https://docs.kedro.org/en/stable/configuration/advanced_configuration.html#how-to-override-configuration-with-runtime-parameters-with-the-omegaconfigloader

[lordsoffallen]

I'm sure i can find a way to hack it but it's kind of a headache to hack things constantly as it suggests it's not designed for those workflows. Again, for standard ML required hacks are minimum so it's fine there but LLM workflows weren't a great fit so far which made to write extra connectors/workarounds and i ended up questioning the value i gained from kedro by using it.

[datajoely]

Kedro design decisions

Kedro was built against scar tissue gained building 'traditional' ML models, let's say sklearn style for the sake of it. It's principles speak about simplicity, modularity and reproducibility. That last point is uncharted territory with non-deterministic systems like LLMs.

One of the tensions in supporting advanced user of Kedro has been the balance between supporting dynamic pipelines. With dynamism one quickly gets questions about conditional logic and our party line has always been some version of -

"we do not believe the combinatorial complexity conditional logic introduced is worth the maintenance burden trade off".

Kedro 🤝 LLM scepticism

Orchestrators have work at a higher level of abstraction than Kedro. This is to say where 1 node is often 1 K8s pod, and have opted to introduce this type of operator where the DAG complexity introduced is lower. Examples (1), (2), (3).

Personally I currently do not believe our current principles or design decisions around Kedro lend themselves well to the requirements of working with LLMs. Whether we change / relax these is another question.... If you look at the effort that's gone into Prefect Control Flow there is a great deal of work that's gone into handling IO of various function calls within the task DAG. This pattern is very different to Kedro's decoupled dataset centric way of doing IO outside of nodes.

Proposition: Let's make Kedro a great fit for deterministic 'function calling'

Our biggest most important issue stems back to #143 and the fact that the Kedro session isn't fit for injecting data at runtime #2169. I believe Kedro could be a great ergonomic solution for LLM systems to perform complex, custom "Function calls" as part of their wider task chains/graphs. If we solve these underlying issues... we'll solve a bunch of problems in #3094 and this LLM space.

[lordsoffallen]

How about keeping standard ML and LLM space separate as a package? kedro does what it does now with focus on standard flow and kedro-llm focuses on improving LLM workflow. I believe trying to optimize for both will not make more people happy but having 2 tools with different designs could potentially work. Not to mention to advantage of not dealing with breaking changes and countless brainstorming about how to integrate LLM workflow into standard kedro flow

[datajoely]

Absolutely open to this- but I think we need to really nail down what kedro-llm does and why it deserves to exist to go in this direction.

[lordsoffallen]

💯. A good approach would be to start understading workflows of LLM case from data scientists - this is important as I would like kedro-llm not become yet another tool made for business and software engineers who has no idea about ML workflow in general. -

[astrojuanlu]

I just sent an email (if you didn't receive it, then I might have sent it to the wrong place 😅)

[ElenaKhaustova]

I see the value in providing a way to use LLMs with Kedro. As a minimum requirement, I would consider providing the mechanism to configure the model, switch between the LLMs, manage prompts, generate output, and get embeddings.

"They really aim for high level users with little to no experience in building ML models."

Here, the question is if "low-lever" users, let's say those who implement models/serve them in production, are going to use Kedro for that. IMO, we can start from the most obvious cases when people just create their pipelines using ready-to-use LLMs/ external APIs and need some mechanism to incorporate them and configure them, as this covers most of the cases.

I would also not avoid using langchain as it provides many things out of the box. Instead, I would try to provide an API on top of it and huggingface so that users could easily switch between different LLMs implementations.

When we worked with the langchain we found it convenient to work with chains - that combine llm and promt and provide a standardised interface to call the model with run-time parameters (prompt placeholders). So one can use different llms with the same interface. Example using latest langchain API: https://python.langchain.com/v0.1/docs/integrations/chat/anthropic/

We also came to dynamic model initialisation, in our case it can help users to switch between different models without need to add extra datasets (OpenAI, Cohere, Azure, etc).

If combine those ideas the catalog.yaml can look like that:

gpt_3_5_turbo:
   type: langchain.LLMDataset
   model_type: langchain_openai.ChatOpenAI
   kwargs:
     model: "gpt-3.5-turbo"
     temperature: 0.0
   credentials: openai
   prompt: "prompt" # we can even use parameters to provide them as it's usually more convenient to keep them in a separate files as they might be too long

         
bert:
   type: huggingface.LLMDataset
   model_type: transformers.TFBertModel
   kwargs:
     model: "bert-base-uncased"
     tokenizer: "bert-base-uncased"

As for dynamic promts - it can be partly solved via placeholders in the prompts that are filled at the runtime as it's done in langchain. So we can inject any content or the entire prompt during execution.

"Kedro was built against scar tissue gained building 'traditional' ML models, let's say sklearn style for the sake of it. It's principles speak about simplicity, modularity and reproducibility. That last point is uncharted territory with non-deterministic systems like LLMs."

I think there's a difference between adding some feature to the framework, which leads to non-reproducible behaviour, and using Kedro with non-deterministic models (via dataset). We cannot control the last as users can make it anyway but should we?