litelines 0.1.1

Customize, control, and enhance LLM generation with logits processors, featuring visualization capabilities to inspect and understand state transitions

litelines

Customize, control, and enhance LLM generation with logits processors, featuring visualization capabilities to inspect and understand state transitions.

Installation

pip install litelines

Dependencies

The only dependency is outlines-core.

Supported Frameworks

• transformers

Basic Usage

• Define a logits processor through a Pydantic schema or a regular expression and visualize it:

from litelines.transformers import SchemaProcessor

processor = SchemaProcessor(r"Yes\.|No\.", tokenizer)
processor.show_graph()

• Generate a structured answer:

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda:0")

messages = [{"role": "user", "content": "Are you sentient?"}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
).to(model.device)
generated = model.generate(**inputs, logits_processor=[processor])
print(tokenizer.decode(generated[0][inputs['input_ids'].shape[-1]:]))
# No.

• Visualize the selected path:

processor.show_graph()

100% Guaranteed Valid JSON answer

• Define a pydantic schema describing the required JSON or provide the JSON schema as a string:

from typing import Literal
from pydantic import BaseModel, Field

class Sentiment(BaseModel):
    """Correctly inferred `Sentiment` with all the required parameters with correct types."""
    label: Literal["positive", "negative"] = Field(
        ..., description="Sentiment of the text"
    )

'''
Alternatively, provide the JSON schema as a sting:
Sentiment = """{'description': 'Correctly inferred `Sentiment` with all the required parameters with correct types.',
 'properties': {'label': {'description': 'Sentiment of the text',
   'enum': ['positive', 'negative'],
   'title': 'Label',
   'type': 'string'}},
 'required': ['label'],
 'title': 'Sentiment',
 'type': 'object'}"""
'''

• Define the processor and visualize it:

from litelines.transformers import SchemaProcessor

processor = SchemaProcessor(Sentiment, tokenizer)
processor.show_graph()

• Generate a structured answer:

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda:0")

messages = [{"role": "user", "content": "What is the sentiment of the following text: 'Awesome'"}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
).to(model.device)
generated = model.generate(**inputs, logits_processor=[processor])
print(tokenizer.decode(generated[0][inputs['input_ids'].shape[-1]:]))
# {"label": "positive"}

• Visualize the selected path:

processor.show_graph()

100% Guaranteed Valid Tool Calling answer

• Define a pydantic schema describing the tool:

from typing import Literal
from pydantic import BaseModel, Field


class Sentiment(BaseModel):
    """Correctly inferred `Sentiment` with all the required parameters with correct types."""
    label: Literal["positive", "negative"] = Field(
        ..., description="Sentiment of the text"
    )

• Define the processor, add the parameter include_tool_call=True and visualize it:

from litelines.transformers import SchemaProcessor

processor = SchemaProcessor(Sentiment, tokenizer, include_tool_call=True)
processor.show_graph()

• Generate a structured answer:

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda:0")

messages = [{"role": "user", "content": "What is the sentiment of the following text: 'Awesome'"}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
).to(model.device)
generated = model.generate(**inputs, logits_processor=[processor])
print(tokenizer.decode(generated[0][inputs['input_ids'].shape[-1]:]))
# <tool_call>
# {"name": "Sentiment", "arguments": {"label": "positive"}}
# </tool_call>

• Visualize the selected path:

processor.show_graph()

Allow Preamble but still get 100% Guaranteed Valid JSON/Tool Calling answer

• Define a pydantic schema describing the required JSON or provide the JSON schema as a string:

from typing import Literal
from pydantic import BaseModel, Field

class Sentiment(BaseModel):
    """Correctly inferred `Sentiment` with all the required parameters with correct types."""
    label: Literal["positive", "negative"] = Field(
        ..., description="Sentiment of the text"
    )

• Define the processor, add parameter allow_preamble=True and visualize it:

from litelines.transformers import SchemaProcessor

processor = SchemaProcessor(Sentiment, tokenizer, allow_preamble=True)
processor.show_graph()

• Generate a structured answer:

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda:0")

messages = [{"role": "user", "content": "What is the sentiment of the following text: 'Awesome'"}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
).to(model.device)
generated = model.generate(**inputs, logits_processor=[processor])
print(tokenizer.decode(generated[0][inputs['input_ids'].shape[-1]:]))
# The sentiment of the text "Awesome" is positive.
# {"label": "positive"}

• Visualize the selected path:

processor.show_graph()