🌸 The open framework for finetuning LLMs on private data
import asyncio

from fastapi import APIRouter

from lib.models.ingest import IngestRequest
from lib.service.flows import create_finetune
from lib.utils.prisma import prisma
from prisma.models import Datasource

router = APIRouter()


@router.post(
    "/ingest",
    name="ingest",
    description="Ingest data",
)
async def ingest(body: IngestRequest):
    """Endpoint for ingesting data"""
    datasource = await prisma.datasource.create(data=body.dict())

    async def run_training_flow(datasource: Datasource):
        try:
            await create_finetune(
                datasource=datasource,
            )
        except Exception as flow_exception:
            raise flow_exception

    asyncio.create_task(run_training_flow(datasource=datasource))
    return {"success": True, "data": datasource}

lib/api/ingest.py
homanp/nagato
from fastapi import APIRouter

router = APIRouter()


@router.post(
    "/invoke",
    name="invoke",
    description="Invoke model",
)
async def invoke():
    """Endpoint for invoking model"""
    return {"success": True, "data": None}

lib/api/invoke.py
homanp/nagato
from fastapi import APIRouter

from lib.utils.prisma import prisma
from prisma import Json

router = APIRouter()


@router.post(
    "/webhook/finetune",
    name="finetune_webhook",
    description="Webhook events for fine-tuning",
)
async def finetune_webhook(body: dict):
    """Endpoint for invoking model"""
    finetune_id = body.get("id")
    status = body.get("status")
    datasources = await prisma.datasource.find_many()
    for ds in datasources:
        if ds.finetune and ds.finetune.get("id") == finetune_id:
            datasource_ = await prisma.datasource.update(
                where={"id": ds.id},
                data={
                    "finetune": Json({**body, **ds.finetune}),
                    "status": "DONE" if status == "succeeded" else "FAILED",
                },
            )
    return {"success": True, "data": datasource_}

lib/api/webhook.py
homanp/nagato
import logging
import time

import colorlog
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware

from lib.routers import router
from lib.utils.prisma import prisma

# Create a color formatter
formatter = colorlog.ColoredFormatter(
    "%(log_color)s%(levelname)s:  %(message)s",
    log_colors={
        "DEBUG": "cyan",
        "INFO": "green",
        "WARNING": "yellow",
        "ERROR": "red",
        "CRITICAL": "bold_red",
    },
    secondary_log_colors={},
    style="%",
)  # Create a console handler and set the formatter
console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)

logging.basicConfig(
    level=logging.INFO,
    format="%(levelname)s: %(message)s",
    handlers=[console_handler],
    force=True,
)

app = FastAPI(
    title="Superagent",
    docs_url="/",
    description="Build, manage and deploy LLM-powered Agents",
    version="0.1.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    print(f"Total request time: {process_time} secs")
    return response


@app.on_event("startup")
async def startup():
    await prisma.connect()


@app.on_event("shutdown")
async def shutdown():
    await prisma.disconnect()


app.include_router(router)

lib/main.py
homanp/nagato
from typing import Optional

from pydantic import BaseModel


class IngestRequest(BaseModel):
    type: str
    base_model: str
    provider: str
    url: Optional[str]
    content: Optional[str]
    webhook_url: Optional[str]

lib/models/ingest.py
homanp/nagato
from fastapi import APIRouter

from lib.api import ingest, invoke, webhook

router = APIRouter()
api_prefix = "/api/v1"

router.include_router(ingest.router, tags=["Ingest"], prefix=api_prefix)
router.include_router(invoke.router, tags=["Invoke"], prefix=api_prefix)
router.include_router(webhook.router, tags=["Webhook"], prefix=api_prefix)

lib/routers.py
homanp/nagato
from tempfile import NamedTemporaryFile
from typing import List, Union

import requests
from decouple import config
from llama_index import Document, SimpleDirectoryReader
from llama_index.node_parser import SimpleNodeParser
from numpy import ndarray
from sentence_transformers import SentenceTransformer

from lib.service.vectordb import get_vector_service
from prisma.models import Datasource


class EmbeddingService:
    def __init__(self, datasource: Datasource):
        self.datasource = datasource

    def get_datasource_suffix(self) -> str:
        suffixes = {"TXT": ".txt", "PDF": ".pdf", "MARKDOWN": ".md"}
        try:
            return suffixes[self.datasource.type]
        except KeyError:
            raise ValueError("Unsupported datasource type")

    async def generate_documents(self) -> List[Document]:
        with NamedTemporaryFile(
            suffix=self.get_datasource_suffix(), delete=True
        ) as temp_file:
            if self.datasource.url:
                content = requests.get(self.datasource.url).content
            else:
                content = self.datasource.content
            temp_file.write(content)
            temp_file.flush()
            reader = SimpleDirectoryReader(input_files=[temp_file.name])
            docs = reader.load_data()
            return docs

    async def generate_chunks(
        self, documents: List[Document]
    ) -> List[Union[Document, None]]:
        parser = SimpleNodeParser.from_defaults(chunk_size=350, chunk_overlap=20)
        nodes = parser.get_nodes_from_documents(documents, show_progress=True)
        return nodes

    async def generate_embeddings(
        self, nodes: List[Union[Document, None]]
    ) -> List[ndarray]:
        vectordb = await get_vector_service(
            provider="pinecone",
            index_name="all-minilm-l6-v2",
            namespace=self.datasource.id,
            dimension=384,
        )
        model = SentenceTransformer(
            "all-MiniLM-L6-v2", use_auth_token=config("HF_API_KEY")
        )
        embeddings = []
        for node in nodes:
            if node is not None:
                embedding = (
                    node.id_,
                    model.encode(node.text).tolist(),
                    {**node.metadata, "content": node.text},
                )
                embeddings.append(embedding)
        await vectordb.upsert(vectors=embeddings)
        return embeddings

    # def generate_query(self):
    #    model = SentenceTransformer(
    #        "all-MiniLM-L6-v2", use_auth_token=config("HF_API_KEY")
    #    )
    #    vectordb = get_vector_service(
    #        provider="pinecone",
    #        index_name="all-minilm-l6-v2",
    #        namespace=self.datasource.id,
    #        dimension=384,
    #    )
    #    query = "How many cars were sold?"
    #    embedding = model.encode([query]).tolist()
    #    return vectordb.query(queries=embedding, top_k=5, include_metadata=True)

lib/service/embedding.py
homanp/nagato
# flake8: noqa

import asyncio
import json
import os
import uuid
from abc import ABC, abstractmethod
from typing import Dict, List, Union

import httpx
import openai
import replicate
from decouple import config
from llama_index import Document

from lib.service.prompts import (
    GPT_DATA_FORMAT,
    REPLICATE_FORMAT,
    generate_qa_pair_prompt,
)

openai.api_key = config("OPENAI_API_KEY")

REPLICATE_MODELS = {
    "LLAMA2_7B_CHAT": "meta/llama-2-7b-chat:8e6975e5ed6174911a6ff3d60540dfd4844201974602551e10e9e87ab143d81e",
    "LLAMA2_7B": "meta/llama-2-7b:527827021d8756c7ab79fde0abbfaac885c37a3ed5fe23c7465093f0878d55ef",
    "LLAMA2_13B_CHAT": "meta/llama-2-13b-chat:f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
    "LLAMA2_13B": "meta/llama-2-13b:078d7a002387bd96d93b0302a4c03b3f15824b63104034bfa943c63a8f208c38",
    "LLAMA2_70B_CHAT": "meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3",
    "LLAMA2_70B": "meta/llama-2-70b:a52e56fee2269a78c9279800ec88898cecb6c8f1df22a6483132bea266648f00",
    "GPT_J_6B": "replicate/gpt-j-6b:b3546aeec6c9891f0dd9929c2d3bedbf013c12e02e7dd0346af09c37e008c827",
    "DOLLY_V2_12B": "replicate/dolly-v2-12b:ef0e1aefc61f8e096ebe4db6b2bacc297daf2ef6899f0f7e001ec445893500e5",
}

OPENAI_MODELS = {"GPT_35_TURBO": "gpt-3.5-turbo"}


class FinetuningService(ABC):
    def __init__(self, nodes: List[Union[Document, None]]):
        self.nodes = nodes

    @abstractmethod
    async def generate_dataset(self) -> str:
        pass

    @abstractmethod
    async def validate_dataset(self, training_file: str) -> str:
        pass

    @abstractmethod
    async def finetune(self, training_file: str, base_model: str) -> Dict:
        pass

    async def cleanup(self, training_file: str) -> None:
        os.remove(training_file)


class OpenAIFinetuningService(FinetuningService):
    def __init__(
        self,
        nodes: List[Union[Document, None]],
        num_questions_per_chunk: int = 10,
        batch_size: int = 10,
        base_model: str = "GPT_35_TURBO",
    ):
        super().__init__(nodes=nodes)
        self.num_questions_per_chunk = num_questions_per_chunk
        self.batch_size = batch_size
        self.base_model = base_model

    async def generate_prompt_and_completion(self, node):
        prompt = generate_qa_pair_prompt(
            context=node.text, num_of_qa_paris=10, format=GPT_DATA_FORMAT
        )
        completion = await openai.ChatCompletion.acreate(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
        )
        return completion.choices[0].message.content

    async def generate_dataset(self) -> str:
        training_file = f"{uuid.uuid4()}.jsonl"
        with open(training_file, "w") as f:
            for i in range(
                0, len(self.nodes), self.batch_size
            ):  # Process nodes in chunks of batch_size
                tasks = [
                    self.generate_prompt_and_completion(node)
                    for node in self.nodes[i : i + self.batch_size]
                ]
                qa_pairs = await asyncio.gather(*tasks)
                for qa_pair in qa_pairs:
                    json_objects = qa_pair.split("\n\n")
                    for json_obj in json_objects:
                        f.write(json_obj + "\n")
        return training_file

    async def validate_dataset(self, training_file: str) -> str:
        pass

    async def finetune(self, training_file: str) -> Dict:
        file = openai.File.create(file=open(training_file, "rb"), purpose="fine-tune")
        finetune = await openai.FineTuningJob.acreate(
            training_file=file.get("id"), model=OPENAI_MODELS[self.base_model]
        )
        return {**finetune, "training_file": training_file}


class ReplicateFinetuningService(FinetuningService):
    def __init__(
        self,
        nodes: List[Union[Document, None]],
        num_questions_per_chunk: int = 1,
        batch_size: int = 10,
        base_model: str = "LLAMA2_7B_CHAT",
    ):
        super().__init__(nodes=nodes)
        self.num_questions_per_chunk = num_questions_per_chunk
        self.batch_size = batch_size
        self.base_model = base_model

    async def generate_prompt_and_completion(self, node):
        prompt = generate_qa_pair_prompt(
            context=node.text, num_of_qa_paris=10, format=REPLICATE_FORMAT
        )
        completion = await openai.ChatCompletion.acreate(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
        )
        return completion.choices[0].message.content

    async def generate_dataset(self) -> str:
        training_file = f"{uuid.uuid4()}.jsonl"
        with open(training_file, "w") as f:
            for i in range(
                0, len(self.nodes), self.batch_size
            ):  # Process nodes in chunks of batch_size
                tasks = [
                    self.generate_prompt_and_completion(node)
                    for node in self.nodes[i : i + self.batch_size]
                ]
                qa_pairs = await asyncio.gather(*tasks)
                for qa_pair in qa_pairs:
                    json_objects = qa_pair.split("\n\n")
                    for json_obj in json_objects:
                        f.write(json_obj + "\n")
        return training_file

    async def validate_dataset(self, training_file: str) -> str:
        valid_data = []
        with open(training_file, "r") as f:
            for line in f:
                data = json.loads(line)
                if "prompt" in data and "completion" in data:
                    valid_data.append(data)
        with open(training_file, "w") as f:
            for data in valid_data:
                f.write(json.dumps(data) + "\n")
        return training_file

    async def finetune(self, training_file: str) -> Dict:
        training_file_url = await upload_replicate_dataset(training_file=training_file)
        training = replicate.trainings.create(
            version=REPLICATE_MODELS[self.base_model],
            input={
                "train_data": training_file_url,
                "num_train_epochs": 6,
            },
            destination="homanp/test",
            webhook="https://api.nagato.sh/api/v1/webhook/finetune",
        )
        return {"id": training.id, "training_file": training_file}


async def get_finetuning_service(
    nodes: List[Union[Document, None]],
    provider: str = "openai",
    base_model: str = "GPT_35_TURBO",
    num_questions_per_chunk: int = 10,
    batch_size: int = 10,
):
    services = {
        "OPENAI": OpenAIFinetuningService,
        "REPLICATE": ReplicateFinetuningService,
        # Add other providers here
    }
    service = services.get(provider)
    if service is None:
        raise ValueError(f"Unsupported provider: {provider}")
    return service(
        nodes=nodes,
        num_questions_per_chunk=num_questions_per_chunk,
        batch_size=batch_size,
        base_model=base_model,
    )


async def upload_replicate_dataset(training_file: str) -> str:
    headers = {"Authorization": f"Token {config('REPLICATE_API_TOKEN')}"}
    async with httpx.AsyncClient() as client:
        response = await client.post(
            "https://dreambooth-api-experimental.replicate.com/v1/upload/data.jsonl",
            headers=headers,
        )
        response_data = response.json()
        upload_url = response_data["upload_url"]

        with open(training_file, "rb") as f:
            await client.put(
                upload_url,
                headers={"Content-Type": "application/jsonl"},
                content=f.read(),
            )

        serving_url = response_data["serving_url"]
        return serving_url

lib/service/finetune.py
homanp/nagato
from typing import List, Union

import openai
from llama_index import Document
from prefect import flow, task

from lib.service.embedding import EmbeddingService
from lib.service.finetune import get_finetuning_service
from lib.utils.prisma import prisma
from prisma import Json
from prisma.models import Datasource


@task
async def create_vector_embeddings(
    datasource: Datasource,
) -> List[Union[Document, None]]:
    embedding_service = EmbeddingService(datasource=datasource)
    documents = await embedding_service.generate_documents()
    nodes = await embedding_service.generate_chunks(documents=documents)
    await embedding_service.generate_embeddings(nodes=nodes)
    return nodes


@task
async def create_finetuned_model(datasource: Datasource):
    embedding_service = EmbeddingService(datasource=datasource)
    documents = await embedding_service.generate_documents()
    nodes = await embedding_service.generate_chunks(documents=documents)
    finetunning_service = await get_finetuning_service(
        nodes=nodes,
        provider=datasource.provider,
        batch_size=5,
        base_model=datasource.base_model,
    )
    training_file = await finetunning_service.generate_dataset()
    formatted_training_file = await finetunning_service.validate_dataset(
        training_file=training_file
    )
    finetune = await finetunning_service.finetune(training_file=formatted_training_file)
    if datasource.provider == "OPENAI":
        finetune = await openai.FineTune.retrieve(id=finetune.get("id"))
    await finetunning_service.cleanup(training_file=finetune.get("training_file"))
    return finetune


@flow(name="create_finetune", description="Create a finetune", retries=0)
async def create_finetune(datasource: Datasource):
    await create_vector_embeddings(datasource=datasource)
    finetune = await create_finetuned_model(datasource=datasource)
    print(finetune)
    await prisma.datasource.update(
        where={"id": datasource.id},
        data={"finetune": Json(data=finetune)},
    )

lib/service/flows.py
homanp/nagato
# flake8: noqa

GPT_DATA_FORMAT = (
    "{"
    '"messages": ['
    '{"role": "system", "content": "You are an AI agent that\'s an expert at answering questions."}, '
    '{"role": "user", "content": "What\'s the capital of France?"}, '
    '{"role": "assistant", "content": "Paris, is the capital of France."}'
    "]"
    "}"
)

REPLICATE_FORMAT = (
    "{"
    '"prompt": "What\'s the capital of France?",'
    '"completion": "Paris, is the capital of France"'
    "}"
)


def generate_qa_pair_prompt(format: str, context: str, num_of_qa_paris: int = 10):
    prompt = (
        "You are an AI assistant tasked with generating question and answer pairs"
        "for the given context using the given format. Only answer in the format with"
        f"no other text. You should create the following number of question/answer pairs: {num_of_qa_paris}"
        "Return the question/answer pairs as a JSONL."
        "Each dict in the list should have the full context provided,"
        "a relevant question to the context and an answer to the question.\n\n"
        f"Format:\n {format}\n\n"
        f"Context:\n {context}"
    )
    return prompt

lib/service/prompts.py
homanp/nagato
from abc import ABC, abstractmethod
from typing import List

import pinecone
from decouple import config
from numpy import ndarray


class VectorDBService(ABC):
    def __init__(self, index_name: str, dimension: int, namespace: str = None):
        self.index_name = index_name
        self.namespace = namespace
        self.dimension = dimension

    @abstractmethod
    def upsert():
        pass

    @abstractmethod
    def query():
        pass


class PineconeVectorService(VectorDBService):
    def __init__(self, index_name: str, dimension: int, namespace: str = None):
        super().__init__(
            index_name=index_name, dimension=dimension, namespace=namespace
        )
        pinecone.init(api_key=config("PINECONE_API_KEY"))
        # Create a new vector index if it doesn't
        # exist dimensions should be passed in the arguments
        if index_name not in pinecone.list_indexes():
            pinecone.create_index(
                name=index_name, metric="cosine", shards=1, dimension=dimension
            )
        self.index = pinecone.Index(index_name=self.index_name)

    async def upsert(self, vectors: ndarray):
        self.index.upsert(vectors=vectors, namespace=self.namespace)

    async def query(
        self, queries: List[ndarray], top_k: int, include_metadata: bool = True
    ):
        return self.index.query(
            queries=queries,
            top_k=top_k,
            include_metadata=include_metadata,
            namespace=self.namespace,
        )


async def get_vector_service(
    provider: str, index_name: str, namespace: str = None, dimension: int = 384
):
    services = {
        "pinecone": PineconeVectorService,
        # Add other providers here
        # e.g "weaviate": WeaviateVectorService,
    }
    service = services.get(provider)
    if service is None:
        raise ValueError(f"Unsupported provider: {provider}")
    return service(index_name=index_name, namespace=namespace, dimension=dimension)

lib/service/vectordb.py
homanp/nagato
from prisma import Prisma

prisma = Prisma()

lib/utils/__init__.py
homanp/nagato
from prisma import Prisma

prisma = Prisma()

lib/utils/prisma.py
homanp/nagato