Merge branch 'generative-computing:main' into main

Author: tuliocoppola

README.md

@@ -47,12 +47,30 @@ You can get started with a local install, or by using Colab notebooks.
 
 <img src="https://github.com/generative-computing/mellea/raw/main/docs/GetStarted_py.png" style="max-width:800px">
 
-Install with pip:
+Install with [uv](https://docs.astral.sh/uv/getting-started/installation/):
 
 ```bash
 uv pip install mellea
 ```
 
+Install with pip:
+
+```bash
+pip install mellea
+```
+
+> [!NOTE]
+> `mellea` comes with some additional packages as defined in our `pyproject.toml`. I you would like to install all the extra optional dependencies, please run the following commands:
+>
+> ```bash
+> uv pip install mellea[hf] # for Huggingface extras and Alora capabilities.
+> uv pip install mellea[watsonx] # for watsonx backend
+> uv pip install mellea[docling] # for docling
+> uv pip install mellea[all] # for all the optional dependencies
+> ```
+>
+> You can also install all the optional dependencies with `uv sync --all-extras`
+
 > [!NOTE]
 > If running on an Intel mac, you may get errors related to torch/torchvision versions. Conda maintains updated versions of these packages. You will need to create a conda environment and run `conda install 'torchvision>=0.22.0'` (this should also install pytorch and torchvision-extra). Then, you should be able to run `uv pip install mellea`. To run the examples, you will need to use `python <filename>` inside the conda environment instead of `uv run --with mellea <filename>`.
 
@@ -110,7 +128,19 @@ uv venv .venv && source .venv/bin/activate
 Use `uv pip` to install from source with the editable flag:
 
 ```bash
-uv pip install -e .
+uv pip install -e .[all]
+```
+
+If you are planning to contribute to the repo, it would be good to have all the development requirements installed:
+
+```bash
+uv pip install .[all] --group dev --group notebook --group docs
+```
+
+or 
+
+```bash
+uv sync --all-extras --all-groups
 ```
 
 Ensure that you install the precommit hooks:

docs/examples/generative_slots/generative_slots.py

@@ -29,6 +29,6 @@ def generate_summary(text: str) -> str:
             surface. Compared with other rays, they have long tails, and well-defined, rhomboidal bodies.
             They are ovoviviparous, giving birth to up to six young at a time. They range from 0.48 to
             5.1 m (1.6 to 16.7 ft) in length and 7 m (23 ft) in wingspan.
-            """,
+            """
         )
         print("Generated summary is :", summary)

docs/examples/instruct_validate_repair/101_email.py

@@ -1,16 +1,16 @@
 # This is the 101 example for using `session` and `instruct`.
 # helper function to wrap text
 from docs.examples.helper import w
-from mellea import start_session, instruct
+from mellea import instruct, start_session
 from mellea.backends.types import ModelOption
 
 # create a session using Granite 3.3 8B on Ollama and a simple context [see below]
 with start_session(model_options={ModelOption.MAX_NEW_TOKENS: 200}):
-# write an email
+    # write an email
     email_v1 = instruct("Write an email to invite all interns to the office party.")
 
 with start_session(model_options={ModelOption.MAX_NEW_TOKENS: 200}) as m:
-# write an email
+    # write an email
     email_v1 = m.instruct("Write an email to invite all interns to the office party.")
 
 # print result

mellea/__init__.py

@@ -3,18 +3,26 @@
 import mellea.backends.model_ids as model_ids
 from mellea.stdlib.base import LinearContext, SimpleContext
 from mellea.stdlib.genslot import generative
-from mellea.stdlib.session import MelleaSession, start_session, instruct, chat, validate, query, transform
+from mellea.stdlib.session import (
+    MelleaSession,
+    chat,
+    instruct,
+    query,
+    start_session,
+    transform,
+    validate,
+)
 
 __all__ = [
     "LinearContext",
     "MelleaSession",
     "SimpleContext",
+    "chat",
     "generative",
+    "instruct",
     "model_ids",
+    "query",
     "start_session",
-    "instruct",
-    "chat",
+    "transform",
     "validate",
-    "query",
-    "transform"
 ]

mellea/backends/huggingface.py

@@ -150,7 +150,7 @@ def __init__(
                 # Get the model and tokenizer.
                 self._model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
                     self._hf_model_id
-                ).to(self._device)
+                ).to(self._device)  # type: ignore
                 self._tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(
                     self._hf_model_id
                 )

mellea/stdlib/genslot.py

@@ -153,7 +153,11 @@ def __init__(self, func: Callable[P, R]):
         functools.update_wrapper(self, func)
 
     def __call__(
-        self, m=None, model_options: dict | None = None, *args: P.args, **kwargs: P.kwargs
+        self,
+        m=None,
+        model_options: dict | None = None,
+        *args: P.args,
+        **kwargs: P.kwargs,
     ) -> R:
         """Call the generative slot.
 

mellea/stdlib/safety/guardian.py

@@ -125,7 +125,7 @@ def _guardian_validate(self, ctx: Context):
         model = AutoModelForCausalLM.from_pretrained(
             self._model_version, device_map="auto", torch_dtype=torch.bfloat16
         )
-        model.to(self._device)
+        model.to(self._device)  # type: ignore
         model.eval()
 
         tokenizer = AutoTokenizer.from_pretrained(self._model_version)

mellea/stdlib/session.py

@@ -3,8 +3,10 @@
 from __future__ import annotations
 
 import contextvars
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any, Generator, Literal, Optional
+from copy import deepcopy
+from typing import Any, Literal, Optional
 
 from mellea.backends import Backend, BaseModelSubclass
 from mellea.backends.formatter import FormatterBackend
@@ -33,14 +35,13 @@
 from mellea.stdlib.requirement import Requirement, ValidationResult, check, req
 from mellea.stdlib.sampling import SamplingResult, SamplingStrategy
 
-
 # Global context variable for the context session
-_context_session: contextvars.ContextVar[Optional["MelleaSession"]] = contextvars.ContextVar(
+_context_session: contextvars.ContextVar[MelleaSession | None] = contextvars.ContextVar(
     "context_session", default=None
 )
 
 
-def get_session() -> "MelleaSession":
+def get_session() -> MelleaSession:
     """Get the current session from context.
 
     Raises:
@@ -71,6 +72,7 @@ def backend_name_to_class(name: str) -> Any:
     else:
         return None
 
+
 def start_session(
     backend_name: Literal["ollama", "hf", "openai", "watsonx"] = "ollama",
     model_id: str | ModelIdentifier = IBM_GRANITE_3_3_8B,
@@ -147,6 +149,7 @@ def start_session(
     backend = backend_class(model_id, model_options=model_options, **backend_kwargs)
     return MelleaSession(backend, ctx)
 
+
 class MelleaSession:
     """Mellea sessions are a THIN wrapper around `m` convenience functions with NO special semantics.
 
@@ -451,13 +454,23 @@ def genslot(
         Returns:
             ModelOutputThunk: Output thunk
         """
+        generate_logs: list[GenerateLog] = []
         result: ModelOutputThunk = self.backend.generate_from_context(
             action=gen_slot,
             ctx=self.ctx,
             model_options=model_options,
             format=format,
+            generate_logs=generate_logs,
             tool_calls=tool_calls,
         )
+        # make sure that the last and only Log is marked as the one related to result
+        assert len(generate_logs) == 1, "Simple call can only add one generate_log"
+        generate_logs[0].is_final_result = True
+
+        self.ctx.insert_turn(
+            ContextTurn(deepcopy(gen_slot), result), generate_logs=generate_logs
+        )
+
         return result
 
     def query(

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "typer",
     "click<8.2.0", # Newer versions will cause errors with --help in typer CLIs.
     "mistletoe>=1.4.0",
+    "huggingface-hub>=0.33.4",
 ]
 
 [project.scripts]
@@ -67,6 +68,8 @@ docling = [
     "docling>=2.45.0",
 ]
 
+all = ["mellea[watsonx,docling,hf]"]
+
 [dependency-groups]
 # Use these like:
 # pip install -e . --group dev

test/stdlib_basics/test_genslot.py

@@ -1,6 +1,7 @@
 import pytest
 from typing import Literal
 from mellea import generative, start_session
+from mellea.stdlib.base import LinearContext
 
 
 @generative
@@ -13,7 +14,7 @@ def write_me_an_email() -> str: ...
 
 @pytest.fixture
 def session():
-    return start_session()
+    return start_session(ctx=LinearContext())
 
 
 @pytest.fixture
@@ -34,5 +35,11 @@ def test_sentiment_output(classify_sentiment_output):
     assert classify_sentiment_output in ["positive", "negative"]
 
 
+def test_gen_slot_logs(classify_sentiment_output, session):
+    sent = classify_sentiment_output
+    last_prompt = session.last_prompt()[-1]
+    assert isinstance(last_prompt, dict)
+    assert set(last_prompt.keys()) == {"role", "content"}
+
 if __name__ == "__main__":
     pytest.main([__file__])