A pythonic library providing light-weighted interface with LLMs
# -*- coding: utf-8 -*-
import requests
from llm_core.splitters import TokenSplitter
from llm_core.assistants import (
    Analyst,
    Doubter,
    ConsistencyVerifier,
    LLaMACPPAssistant,
)

# Fetch some content
pizza_dough_recipe_url = (
    "https://raw.githubusercontent.com/hendricius/pizza-dough/main/README.md"
)

model = "mistral-7b-instruct-v0.1.Q4_K_M.gguf"
assistant_cls = LLaMACPPAssistant

# Utilities
analyst = Analyst(model, assistant_cls)
doubter = Doubter(model, assistant_cls)
verifier = ConsistencyVerifier(model, assistant_cls)

splitter = TokenSplitter(model=model, chunk_size=3_000)
pizza_dough_recipe = requests.get(pizza_dough_recipe_url).text
context = splitter.first_extract(pizza_dough_recipe)

query = "Write 3 advices when making pizza dough."


analyst_response = analyst.ask(query, context)
print(analyst_response.content)

question_collection = doubter.verify(query, analyst_response.content)
questions = question_collection.questions
print(questions)


answers = []

for question in questions:
    response = analyst.ask(question, context=context)
    answers.append(response.content)

for question, answer in zip(questions, answers):
    verifications = verifier.verify(
        question=question, context=context, answer=response.content
    )
    print(question, answer, verifications)

examples/chain_of_verification.py
paschembri/py-llm-core
import json
from typing import List
from enum import Enum
from dataclasses import dataclass, asdict
from llm_core.assistants import LLaMACPPAssistant


class CRUDOperation(Enum):
    CREATE = 0
    READ = 1
    UPDATE = 2
    DELETE = 3


@dataclass
class UserQueryGenerator:
    system_prompt = "You are a helpful assistant."
    prompt = """
    # Goals

    We are developing a new b2b software able to understand user's queries
    written in plain english.

    # Examples of queries

    Cancel all my meetings of the week
    What is on the agenda for the meeting at 1 pm ?
    Show me the log entries for the last 24h
    Disable all access for john.appleseed@gmail.com
    {examples}

    # Todo

    Write {queries_count} new examples of what users could ask.

    """
    user_queries: List[str]

    @classmethod
    def generate(cls, queries_count=10, examples=()):
        with LLaMACPPAssistant(
            cls,
            model="mistral",
            completion_kwargs={"temperature": 0.8},
        ) as assistant:
            batch = assistant.process(
                queries_count=queries_count, examples=examples
            )
            return batch.user_queries


@dataclass
class UserQueryGeneratorV2:
    system_prompt = "You are a helpful assistant."
    prompt = """
    # Goals

    We are developing a new business software able to understand user's commands
    written in plain english. Then we'll map these commands to simple operations
    like create, read, update and delete.

    # Todo

    Imagine {queries_count} examples of what users could ask.

    """
    user_prompts: List[str]

    @classmethod
    def generate(cls, queries_count=20, examples=()):
        with LLaMACPPAssistant(
            cls,
            model="mistral",
            completion_kwargs={"temperature": 1.0},
        ) as assistant:
            batch = assistant.process(queries_count=queries_count)
            return batch.user_prompts


@dataclass
class UserQueryClassification:
    system_prompt = "You are a helpful assistant."
    prompt = """
    # Instructions

    Analyze the user's query and classify his intent into:

    - an operation (among CRUD)
    - a target entity

    # Examples

    query: Cancel all my meetings of the week
    operation: DELETE
    item: MEETING

    query: Show me the log entries for the last 24h
    operation: READ
    item: LOG

    # Todo

    Classify the query: {prompt}
    """
    operation: CRUDOperation
    item: str

    def to_json(self):
        attrs = asdict(self)
        attrs["operation"] = attrs["operation"].name
        return json.dumps(attrs)

    @classmethod
    def classify(cls, prompt):
        with LLaMACPPAssistant(cls, model="mistral") as assistant:
            user_query = assistant.process(prompt=prompt)
            return user_query


def main():
    dataset_queries = []

    with open("dataset.txt", "a") as file:
        for i in range(5):
            queries = UserQueryGeneratorV2.generate()
            for query in queries:
                print(query, file=file, flush=True)
                dataset_queries.append(query)

    with open("dataset.jsonl", "a") as file:
        for query in dataset_queries:
            classification = UserQueryClassification.classify(query)
            print(classification.to_json(), file=file, flush=True)

examples/synthetic-dataset-generation-example.py
paschembri/py-llm-core
# coding: utf-8
from pathlib import Path
from setuptools import setup, find_packages


here = Path(__file__).parent
packages = find_packages("src")
main_package = packages[0]
long_description = (here / "README.md").read_text()

setup(
    name="py-llm-core",
    version="2.2.0",
    license="MIT",
    description="PyLLMCore provides a light-weighted interface with LLMs",
    long_description=long_description,
    long_description_content_type="text/markdown",
    author="P.A. SCHEMBRI",
    author_email="pa.schembri@advanced-stack.com",
    url="https://github.com/paschembri/py-llm-core",
    packages=packages,
    package_dir={"": "src"},
    include_package_data=True,
    install_requires=[
        "openai",
        "tiktoken",
        "jsonschema",
        "dirtyjson",
        "llama-cpp-python",
        "python-decouple",
    ],
)

setup.py
paschembri/py-llm-core
import codecs

from .token_codecs import lookup

codecs.register(lookup)

src/llm_core/__init__.py
paschembri/py-llm-core
from .base import OpenAIAssistant, LLaMACPPAssistant
from .summarizers import Summarizer, SimpleSummary, DenserSummaryCollection
from .verifiers import (
    QuestionCollection,
    Doubter,
    AnswerConsistency,
    ConsistencyVerifier,
)
from .analysts import Answer, Analyst


__all__ = [
    "OpenAIAssistant",
    "LLaMACPPAssistant",
    "Summarizer",
    "SimpleSummary",
    "DenserSummaryCollection",
    "QuestionCollection",
    "Doubter",
    "AnswerConsistency",
    "ConsistencyVerifier",
    "Answer",
    "Analyst",
]

src/llm_core/assistants/__init__.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
from dataclasses import dataclass


@dataclass
class Answer:
    system_prompt = "You are a helpful analyst."

    prompt = """Context:
    ```
    {context}
    ```

    Using *only* the Context, answer briefly to the following:
    ```
    {question}
    ```
    """

    content: str


@dataclass
class Analyst:
    model: str
    assistant_cls: type
    results_cls: type = Answer

    def batch_ask(self, questions, context):
        for question in questions:
            with self.assistant_cls(
                self.results_cls, model=self.model
            ) as assistant:
                answer = assistant.process(question=question, context=context)

            # Release before yielding
            yield answer

    def ask(self, question, context):
        with self.assistant_cls(
            self.results_cls, model=self.model
        ) as assistant:
            answer = assistant.process(question=question, context=context)
            return answer

src/llm_core/assistants/analysts.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
from ..parsers import BaseParser, OpenAIParser, LLaMACPPParser


class BaseAssistant(BaseParser):
    def __init__(self, target_cls, *args, **kwargs):
        super().__init__(target_cls, *args, **kwargs)
        self.system_prompt = getattr(self.target_cls, "system_prompt", "")
        self.prompt = getattr(self.target_cls, "prompt", "")

    def process(self, **kwargs):
        system_prompt = self.system_prompt.format(**kwargs)
        prompt = self.prompt.format(**kwargs)

        self.model_wrapper.system_prompt = system_prompt

        completion = self.model_wrapper.ask(
            prompt, schema=self.target_json_schema, **self.completion_kwargs
        )
        instance = self.deserialize(completion.choices[0].message.content)
        return instance


class OpenAIAssistant(BaseAssistant, OpenAIParser):
    def __init__(self, target_cls, model="gpt-3.5-turbo", *args, **kwargs):
        super().__init__(target_cls, model=model, *args, **kwargs)


class LLaMACPPAssistant(BaseAssistant, LLaMACPPParser):
    def __init__(
        self,
        target_cls,
        model="mistral",
        llama_cpp_kwargs=None,
        *args,
        **kwargs
    ):
        super().__init__(
            target_cls,
            model=model,
            llama_cpp_kwargs=llama_cpp_kwargs,
            *args,
            **kwargs
        )

src/llm_core/assistants/base.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
from typing import List
from dataclasses import dataclass
from ..splitters import TokenSplitter


@dataclass
class Summary:
    content: str


@dataclass
class SimpleSummary(Summary):
    system_prompt = """
    You are an expert at extracting key information and summarizing.
    """

    prompt = """Content:
    ```
    {content}
    ```
    --
    Summarize the content in {word_count}.
    """


@dataclass
class DenseSummary(Summary):
    missing_entities: List[str]


@dataclass
class DenserSummaryCollection:
    system_prompt = """
    You are an expert in writing rich and dense summaries in broad domains.
    """

    prompt = """Content:
    ```
    {content}
    ```
    --

    You will generate increasingly concise, entity-dense summaries of the above
    Article.

    Repeat the following 2 steps 5 times.

    - Step 1: Identify 1-3 informative Entities from the Article
    which are missing from the previously generated summary and are the most
    relevant.

    - Step 2: Write a new, denser summary of identical length which covers
    every entity and detail from the previous summary plus the missing entities

    A Missing Entity is:

    - Relevant: to the main story
    - Specific: descriptive yet concise (5 words or fewer)
    - Novel: not in the previous summary
    - Faithful: present in the Article
    - Anywhere: located anywhere in the Article

    Guidelines:
    - The first summary should be long (4-5 sentences, approx. {word_count} words) yet
    highly non-specific, containing little information beyond the entities
    marked as missing.

    - Use overly verbose language and fillers (e.g. "this article discusses") to
    reach approx. {word_count} words.

    - Make every word count: re-write the previous summary to improve flow and
    make space for additional entities.

    - Make space with fusion, compression, and removal of uninformative phrases
    like "the article discusses"

    - The summaries should become highly dense and concise yet self-contained,
    e.g., easily understood without the Article.

    - Missing entities can appear anywhere in the new summary.

    - Never drop entities from the previous summary. If space cannot be made,
    add fewer new entities.

    > Remember to use the exact same number of words for each summary.
    Answer in JSON.

    > The JSON in `summaries` should be a list (length 5) of
    dictionaries whose keys are "missing_entities" and "content".

    """

    summaries: List[DenseSummary]


@dataclass
class Summarizer:
    model: str
    assistant_cls: type
    results_cls: type = SimpleSummary

    def fast_summarize(self, content, word_count=80):
        with self.assistant_cls(
            self.results_cls, model=self.model
        ) as assistant:
            splitter = TokenSplitter(
                model=assistant.model_name,
                chunk_size=int(assistant.ctx_size * 0.6),
                chunk_overlap=int(assistant.ctx_size * 0.05),
            )
            chunk = next(splitter.chunkify(content))
            summary = assistant.process(content=chunk, word_count=word_count)
            return summary

    def summarize(self, content, word_count=80):
        with self.assistant_cls(
            self.results_cls, model=self.model
        ) as assistant:
            splitter = TokenSplitter(
                model=assistant.model_name,
                chunk_size=int(assistant.ctx_size * 0.6),
                chunk_overlap=int(assistant.ctx_size * 0.05),
            )
            for chunk in splitter.chunkify(content):
                summary = assistant.process(
                    content=chunk, word_count=word_count
                )
                yield summary

src/llm_core/assistants/summarizers.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
from typing import List
from dataclasses import dataclass


@dataclass
class QuestionCollection:
    system_prompt = "You are a helpful assistant."

    prompt = """Instructions:
    ```
    {instructions}
    ```

    The provided answer was:
    ```
    {answer}
    ```

    Write {n_questions} single-fact, close-ended questions to help verify
    if there are mistakes in the answer.
    """

    questions: List[str]


@dataclass
class Doubter:
    model: str
    assistant_cls: type
    results_cls: type = QuestionCollection

    def verify(self, instructions, answer, n_questions=5):
        with self.assistant_cls(
            self.results_cls, model=self.model
        ) as assistant:
            verification = assistant.process(
                instructions=instructions,
                answer=answer,
                n_questions=n_questions,
            )
            return verification


@dataclass
class AnswerConsistency:
    system_prompt = (
        "You are a meticulous assistant checking for inconsistencies"
    )

    prompt = """Context:
    ```
    {context}
    ```

    Question:
    ```
    {question}
    ```

    The provided Answer was:
    ```
    {answer}
    ```

    - Does the Answer is consistent ?
    - Does the Answer is inferred from the Context ?
    """

    is_consistent: bool
    is_inferred_from_context: bool


@dataclass
class ConsistencyVerifier:
    model: str
    assistant_cls: type
    results_cls: type = AnswerConsistency

    def verify(self, question, context, answer):
        with self.assistant_cls(
            self.results_cls, model=self.model
        ) as assistant:
            verification = assistant.process(
                question=question,
                context=context,
                answer=answer,
            )
            return verification

src/llm_core/assistants/verifiers.py
paschembri/py-llm-core
from .openai import OpenAIChatModel
from .llama_cpp_compatible import LLaMACPPModel

__all__ = ["OpenAIChatModel", "LLaMACPPModel"]

src/llm_core/llm/__init__.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
from dataclasses import dataclass
from typing import List
import codecs
import json


@dataclass
class LLMBase:
    def sanitize_prompt(
        self,
        prompt,
        history=None,
        schema=None,
    ):
        schema_prompt = ""

        if schema:
            schema_prompt = json.dumps(schema_prompt)

        complete_prompt = [
            self.system_prompt,
            prompt,
            schema_prompt,
        ]

        complete_prompt = "\n".join(complete_prompt)

        required_ctx_size = len(codecs.encode(complete_prompt, self.name))
        if required_ctx_size > self.ctx_size:
            raise OverflowError(
                f"Prompt too large {required_ctx_size} for this model {self.ctx_size}"
            )

        return self.ctx_size - required_ctx_size


@dataclass
class Usage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


@dataclass
class Message:
    role: str
    content: str


@dataclass
class ChatCompletionChoice:
    index: int
    message: Message
    finish_reason: str

    @classmethod
    def from_iterable(cls, iterable):
        for item in iterable:
            message_attrs = item["message"]
            if "function_call" in message_attrs:
                message_attrs["content"] = message_attrs["function_call"][
                    "arguments"
                ]
                message_attrs.pop("function_call")

            message = Message(**item["message"])
            index = item["index"]
            finish_reason = item["finish_reason"]
            yield cls(
                index=index, message=message, finish_reason=finish_reason
            )


@dataclass
class ChatCompletion:
    id: str
    object: str
    created: int
    model: str
    choices: List[ChatCompletionChoice]
    usage: Usage

    @classmethod
    def parse(cls, attrs):
        attributes = {}
        attributes.update(attrs)

        attributes["choices"] = list(
            ChatCompletionChoice.from_iterable(attributes["choices"])
        )
        attributes["usage"] = Usage(**attributes["usage"])

        return cls(**attributes)

src/llm_core/llm/base.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
import os
import platform
import llama_cpp

from dataclasses import dataclass

from .base import (
    LLMBase,
    ChatCompletion,
)
from ..schema import to_grammar
from ..settings import MODELS_CACHE_DIR


@dataclass
class LLaMACPPModel(LLMBase):
    """
    Model wrapper using the https://github.com/ggerganov/llama.cpp library

    Not limited to LLaMA models. For the complete list of supported models,

    see https://github.com/ggerganov/llama.cpp#readme
    """

    name: str = "mistral"
    system_prompt: str = "You are a helpful assistant"
    ctx_size: int = 4000
    verbose: bool = False
    llama_cpp_kwargs: dict = None

    def __enter__(self):
        self.load_model()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.release_model()

    def release_model(self):
        del self.model

    def load_model(self):
        if self.llama_cpp_kwargs is None:
            self.llama_cpp_kwargs = {
                "n_ctx": self.ctx_size,
                "verbose": self.verbose,
            }

            if platform.system() == "Darwin" and platform.machine() == "arm64":
                # Offload everything onto the GPU on MacOS
                self.llama_cpp_kwargs["n_gpu_layers"] = 100
                self.llama_cpp_kwargs["n_threads"] = 1

        model_path = os.path.join(MODELS_CACHE_DIR, self.name)
        self.model = llama_cpp.Llama(model_path, **self.llama_cpp_kwargs)

    def ask(
        self,
        prompt,
        history=None,
        schema=None,
        temperature=0,
        **llama_cpp_kwargs,
    ):
        self.sanitize_prompt(prompt=prompt, history=history, schema=schema)

        messages = [
            {
                "role": "system",
                "content": self.system_prompt,
            },
        ]
        if history:
            messages += history

        messages.append(
            {
                "role": "user",
                "content": prompt,
            },
        )

        kwargs = {}
        if schema:
            grammar = to_grammar(schema)
            kwargs = {"grammar": grammar}

        # Allow to call `ask` and free up memory immediately
        model = getattr(self, "model", None)

        if model:
            completion = model.create_chat_completion(
                messages=messages,
                temperature=temperature,
                max_tokens=-1,
                **kwargs,
            )
        else:
            with self as wrapper:
                completion = wrapper.model.create_chat_completion(
                    messages=messages,
                    temperature=temperature,
                    **kwargs,
                )

        return ChatCompletion.parse(completion)

src/llm_core/llm/llama_cpp_compatible.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
import openai

from dataclasses import dataclass

from .base import (
    LLMBase,
    ChatCompletion,
)


@dataclass
class OpenAIChatModel(LLMBase):
    name: str = "gpt-3.5-turbo"
    system_prompt: str = "You are a helpful assistant"

    def __enter__(self):
        # No special initialization required as we are using API
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # No special cleanup required as we are using API
        pass

    @property
    def ctx_size(self):
        if self.name == "gpt-3.5-turbo":
            return 4_000
        elif self.name == "gpt-3.5-turbo-16k":
            return 16_000
        elif self.name == "gpt-4":
            return 8_000
        elif self.name == "gpt-4-32k":
            return 32_000
        else:
            raise KeyError("Unsupported model")

    def ask(
        self,
        prompt,
        history=None,
        schema=None,
        temperature=0,
    ):
        self.sanitize_prompt(
            prompt=prompt,
            history=history,
            schema=schema,
        )

        messages = [
            {
                "role": "system",
                "content": self.system_prompt,
            },
        ]
        if history:
            messages += history

        messages.append(
            {
                "role": "user",
                "content": prompt,
            },
        )

        kwargs = {}
        if schema:
            functions = {
                "name": "PublishAnswer",
                "description": "Publish the answer",
                "parameters": schema,
            }
            function_call = {"name": "PublishAnswer"}

            kwargs = {
                "functions": [functions],
                "function_call": function_call,
            }
        completion = openai.ChatCompletion.create(
            model=self.name,
            messages=messages,
            temperature=temperature,
            **kwargs,
        )
        return ChatCompletion.parse(completion.to_dict_recursive())

src/llm_core/llm/openai.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
import dirtyjson

from .llm import OpenAIChatModel, LLaMACPPModel
from .schema import to_json_schema, from_dict


class BaseParser:
    def __init__(self, target_cls, *args, **kwargs):
        self.target_cls = target_cls
        self.target_json_schema = to_json_schema(self.target_cls)

    def deserialize(self, json_str):
        attributes = dirtyjson.loads(json_str)
        return from_dict(self.target_cls, attributes)

    def parse(self, text):
        prompt = "Extract and parse information from provided content"
        history = [{"role": "user", "content": text}]
        completion = self.model_wrapper.ask(
            prompt, history, schema=self.target_json_schema
        )
        instance = self.deserialize(completion.choices[0].message.content)
        return instance

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


class OpenAIParser(BaseParser):
    def __init__(
        self,
        target_cls,
        model="gpt-3.5-turbo",
        completion_kwargs=None,
        *args,
        **kwargs
    ):
        super().__init__(target_cls, *args, **kwargs)
        self.completion_kwargs = (
            {} if completion_kwargs is None else completion_kwargs
        )

        self.model_wrapper = OpenAIChatModel(
            name=model,
            system_prompt=(
                "Act as a powerful AI able to extract, parse and process "
                "information from unstructured content."
            ),
        )
        self.ctx_size = self.model_wrapper.ctx_size
        self.model_name = self.model_wrapper.name


class LLaMACPPParser(BaseParser):
    """
    Parser using the https://github.com/ggerganov/llama.cpp library

    Not limited to LLaMA models. For the complete list of supported models,

    see https://github.com/ggerganov/llama.cpp#readme
    """

    def __init__(
        self,
        target_cls,
        model="mistral",
        completion_kwargs=None,
        llama_cpp_kwargs=None,
        *args,
        **kwargs
    ):
        super().__init__(target_cls, *args, **kwargs)
        self.completion_kwargs = (
            {} if completion_kwargs is None else completion_kwargs
        )

        self.model_wrapper = LLaMACPPModel(
            name=model, llama_cpp_kwargs=llama_cpp_kwargs
        )
        self.ctx_size = self.model_wrapper.ctx_size
        self.model_name = self.model_wrapper.name

    def __enter__(self):
        self.model_wrapper.load_model()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.model_wrapper.release_model()

src/llm_core/parsers.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
import re
import json
import dataclasses
import typing
from enum import Enum

from llama_cpp.llama_grammar import LlamaGrammar


def to_json_schema(datacls):
    def get_type(field_type):
        if field_type == int:
            return {"type": "integer"}
        elif field_type == str:
            return {"type": "string"}
        elif field_type == bool:
            return {"type": "boolean"}
        elif field_type == float:
            return {"type": "number"}
        elif field_type == complex:
            return {"type": "string", "format": "complex-number"}
        elif field_type == bytes:
            return {"type": "string", "contentEncoding": "base64"}
        elif field_type == tuple:
            return {"type": "array", "items": {}}
        elif field_type == set:
            return {"type": "array", "uniqueItems": True, "items": {}}
        elif isinstance(field_type, typing._GenericAlias):
            if field_type._name == "List":
                return {
                    "type": "array",
                    "items": get_type(field_type.__args__[0]),
                }
            elif field_type._name == "Dict":
                return {
                    "type": "object",
                    "additionalProperties": get_type(field_type.__args__[1]),
                }
        elif dataclasses.is_dataclass(field_type):
            return to_json_schema(field_type)
        elif issubclass(field_type, Enum):
            return {
                "type": "string",
                "enum": list(field_type.__members__.keys()),
            }
        else:
            return {"type": "object"}

    properties = {}
    required = []
    for field in dataclasses.fields(datacls):
        properties[field.name] = get_type(field.type)
        if (
            field.default == dataclasses.MISSING
            and field.default_factory == dataclasses.MISSING
        ):
            required.append(field.name)

    return {"type": "object", "properties": properties, "required": required}


def from_dict(cls, data):
    if dataclasses.is_dataclass(cls):
        field_types = {f.name: f.type for f in dataclasses.fields(cls)}
        return cls(
            **{k: from_dict(field_types[k], v) for k, v in data.items()}
        )
    elif isinstance(cls, typing._GenericAlias):
        if cls._name == "List":
            return [from_dict(cls.__args__[0], v) for v in data]
        elif cls._name == "Dict":
            return {k: from_dict(cls.__args__[1], v) for k, v in data.items()}
    elif issubclass(cls, Enum):
        return getattr(cls, data)
    else:
        return data


SPACE_RULE = '" "?'

PRIMITIVE_RULES = {
    "boolean": '("true" | "false") space',
    "number": '("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? space',
    "integer": '("-"? ([0-9] | [1-9] [0-9]*)) space',
    "string": r""" "\"" (
        [^"\\] |
        "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
      )* "\"" space """,
    "null": '"null" space',
}

INVALID_RULE_CHARS_RE = re.compile(r"[^a-zA-Z0-9-]+")
GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"]')
GRAMMAR_LITERAL_ESCAPES = {"\r": "\\r", "\n": "\\n", '"': '\\"'}


class SchemaConverter:
    def __init__(self, prop_order):
        self._prop_order = prop_order
        self._rules = {"space": SPACE_RULE}

    def _format_literal(self, literal):
        escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub(
            lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)),
            json.dumps(literal),
        )
        return f'"{escaped}"'

    def _add_rule(self, name, rule):
        esc_name = INVALID_RULE_CHARS_RE.sub("-", name)
        if esc_name not in self._rules or self._rules[esc_name] == rule:
            key = esc_name
        else:
            i = 0
            while f"{esc_name}{i}" in self._rules:
                i += 1
            key = f"{esc_name}{i}"
        self._rules[key] = rule
        return key

    def visit(self, schema, name):
        schema_type = schema.get("type")
        rule_name = name or "root"

        if "oneOf" in schema or "anyOf" in schema:
            rule = " | ".join(
                (
                    self.visit(alt_schema, f'{name}{"-" if name else ""}{i}')
                    for i, alt_schema in enumerate(
                        schema.get("oneOf") or schema["anyOf"]
                    )
                )
            )
            return self._add_rule(rule_name, rule)

        elif "const" in schema:
            return self._add_rule(
                rule_name, self._format_literal(schema["const"])
            )

        elif "enum" in schema:
            rule = " | ".join(
                (self._format_literal(v) for v in schema["enum"])
            )
            return self._add_rule(rule_name, rule)

        elif schema_type == "object" and "properties" in schema:
            # TODO: `required` keyword
            prop_order = self._prop_order
            prop_pairs = sorted(
                schema["properties"].items(),
                # sort by position in prop_order (if specified) then by key
                key=lambda kv: (prop_order.get(kv[0], len(prop_order)), kv[0]),
            )

            rule = '"{" space'
            for i, (prop_name, prop_schema) in enumerate(prop_pairs):
                prop_rule_name = self.visit(
                    prop_schema, f'{name}{"-" if name else ""}{prop_name}'
                )
                if i > 0:
                    rule += ' "," space'
                rule += rf' {self._format_literal(prop_name)} space ":" space {prop_rule_name}'
            rule += ' "}" space'

            return self._add_rule(rule_name, rule)

        elif schema_type == "array" and "items" in schema:
            # TODO `prefixItems` keyword
            item_rule_name = self.visit(
                schema["items"], f'{name}{"-" if name else ""}item'
            )
            rule = f'"[" space ({item_rule_name} ("," space {item_rule_name})*)? "]" space'
            return self._add_rule(rule_name, rule)

        else:
            assert (
                schema_type in PRIMITIVE_RULES
            ), f"Unrecognized schema: {schema}"
            return self._add_rule(
                "root" if rule_name == "root" else schema_type,
                PRIMITIVE_RULES[schema_type],
            )

    def format_grammar(self):
        return "\n".join(
            (f"{name} ::= {rule}" for name, rule in self._rules.items())
        )


def to_grammar(schema):
    converter = SchemaConverter({})
    converter.visit(schema, "")
    grammar_string = converter.format_grammar()
    return LlamaGrammar.from_string(grammar_string, verbose=False)

src/llm_core/schema.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
import os
from decouple import config

MODELS_CACHE_DIR = config(
    "MODELS_CACHE_DIR",
    default=os.path.expanduser("~/.cache/py-llm-core/models"),
)


if not os.path.exists(MODELS_CACHE_DIR):
    os.makedirs(MODELS_CACHE_DIR)

# Store models in the same directory to simplify purges
os.environ["TIKTOKEN_CACHE_DIR"] = MODELS_CACHE_DIR

src/llm_core/settings.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
import itertools
import codecs

from dataclasses import dataclass


@dataclass
class TokenSplitter:
    model: str = "gpt-3.5-turbo"
    chunk_size: int = 300
    chunk_overlap: int = 0

    def __post_init__(self):
        if self.chunk_overlap > self.chunk_size:
            raise IndexError(
                f"Overlap {self.chunk_overlap} > window size {self.chunk_size}"
            )

    def first_extract(self, text: str):
        return next(self.chunkify(text))

    def compute_token_count(self, text: str):
        return len(codecs.encode(text, self.model))

    def chunkify(self, text: str):
        tokens = codecs.encode(text, self.model)

        start = 0
        length = len(tokens)

        while True:
            stop = start + self.chunk_size
            chunk = list(itertools.islice(tokens, start, stop))
            yield codecs.decode(chunk, self.model)

            if stop >= length:
                break

            start += self.chunk_size - self.chunk_overlap

src/llm_core/splitters.py
paschembri/py-llm-core
# -*- coding: utf-8 -*-
import codecs
import tiktoken
import llama_cpp
import os

from .llm import LLaMACPPModel

from .settings import MODELS_CACHE_DIR

codec_to_model_name = {}
for model_name in os.listdir(MODELS_CACHE_DIR):
    key = model_name.replace("-", "_").lower()
    codec_to_model_name[key] = model_name


class BaseTokenCodec:
    def encode(self, input, errors="strict"):
        output = self.token_encoder.encode(input)
        return output, len(output)

    def decode(self, input, errors="strict"):
        output = self.token_encoder.decode(input)
        return output, len(output)


class OpenAIGPTCodec(BaseTokenCodec):
    @classmethod
    def is_supported(cls, codec_name):
        return codec_name in (
            "gpt3.5",
            "gpt_3.5",
            "gpt_3.5_turbo",
            "gpt_3.5_turbo_16k",
            "gpt_4",
            "gpt_4_32k",
        )

    def __init__(self, codec_name):
        # Map everything to gpt-3.5-turbo
        self.token_encoder = tiktoken.encoding_for_model("gpt-3.5-turbo")


class LLaMACodec:
    @classmethod
    def is_supported(cls, codec_name):
        return codec_name in codec_to_model_name

    def __init__(self, codec_name):
        self.model_name = codec_to_model_name[codec_name]

    def encode(self, input, errors="strict"):
        with LLaMACPPModel(name=self.model_name) as wrapper:
            tokenizer = llama_cpp.LlamaTokenizer(wrapper.model)
            output = tokenizer.encode(input)

        return output, len(output)

    def decode(self, input, errors="strict"):
        with LLaMACPPModel(name=self.model_name) as wrapper:
            tokenizer = llama_cpp.LlamaTokenizer(wrapper.model)
            output = tokenizer.decode(input)

        return output, len(output)


def lookup(codec_name):
    if OpenAIGPTCodec.is_supported(codec_name):
        codec = OpenAIGPTCodec(codec_name)

    elif LLaMACodec.is_supported(codec_name):
        codec = LLaMACodec(codec_name)

    else:
        return None

    return codecs.CodecInfo(
        name=codec_name,
        encode=codec.encode,
        decode=codec.decode,
    )

src/llm_core/token_codecs.py
paschembri/py-llm-core