Function rag_server

generation/utilities/rag_server/rag_server.py:629–1009

(
    prompt_path,
    template_path,
    gguf_model_path,
    context_length,
    documents_dir: str,
    questions_jsonl_path: str,
    question_chunk_size: int,
    top_k: int,
    cache_dir: str,
    max_shrink_iterations,
    collection_name: str = "questions_collection",
    llama_path="./llama.cpp",  # customizable llama.cpp path
    port=8003,
    task_id=None,
    **kwargs,
)

async def rag_server(
    prompt_path,
    template_path,
    gguf_model_path,
    context_length,
    documents_dir: str,
    questions_jsonl_path: str,
    question_chunk_size: int,
    top_k: int,
    cache_dir: str,
    max_shrink_iterations,
    collection_name: str = "questions_collection",
    llama_path="./llama.cpp",  # customizable llama.cpp path
    port=8003,
    task_id=None,
    **kwargs,
):

    with open(prompt_path, "r") as f:
        prompt = f.read()

    with open(template_path, "r") as f:
        template = f.read()

    # Get parent directory of gguf model path
    model_dir = os.path.dirname(gguf_model_path)

    # Initialize model-specific token counting
    try:
        count_tokens = count_tokens_specific_model(
            model_dir
        )  # Assumes that tokenizer and sucha re still inside the same dir as the saved gguf model.
    except Exception as e:
        print(e)
        print(
            "\n\nYou probably deleted the tokenizer and other ssuch things from the model directory after you got your quantized model. However, the model tokenizer is used to count tokens before sending it off to llama.cpp, so you shouldn't have done that. You can delete model files (*.safetensors) to save space, but leave the tokenizer alone."
        )
        print(
            "To fix this, re-run the datagen pipeline that produced this model. It won't re-train, but it will re-download"
        )
        raise

    # First thing's first, ragify the docs

    app = FastAPI(
        title="Augmentoolkit Custom Model RAG-Enabled API Server", version="0.1.0"
    )

    chroma_client = None
    collection = None
    bm25_index = None
    bm25_corpus_data = []

    # BM25 Persistence paths
    BM25_CACHE_DIR = os.path.join(cache_dir, "bm25_cache")
    BM25_INDEX_PATH = os.path.join(BM25_CACHE_DIR, "bm25_index.pkl")
    BM25_CORPUS_DATA_PATH = os.path.join(BM25_CACHE_DIR, "bm25_corpus_data.json")