File size: 8,016 Bytes
59e6760 4f9b2d4 a529c4d 59e6760 4f9b2d4 a529c4d 4f9b2d4 59e6760 4f9b2d4 a529c4d 59e6760 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field
# Models
class OptimizeRequest(BaseModel):
"""
🔧 Explicit optimization request for RAG (Retrieval-Augmented Generation) pipelines.
Args:
docs_path (str, optional): 📂 Folder containing your documents for RAG optimization. Default: "data/docs"
retriever (List[str], optional): 🔍 Retriever type(s) to use. Default: ['faiss']. Example: 'bm25', 'faiss', 'chroma'
embedding_model (List[str], optional): 🧠 Embedding model(s) to use. Default: ['sentence-transformers/all-MiniLM-L6-v2']
strategy (List[str], optional): 🎯 RAG strategy to apply. Default: ['fixed']. Options: 'fixed', 'token', 'sentence'
chunk_sizes (List[int], optional): 📏 List of chunk sizes to evaluate. Default: [200, 400, 600]
overlaps (List[int], optional): 🔁 List of overlap values to test. Default: [50, 100, 200]
rerankers (List[str], optional): ⚖️ Rerankers to apply after retrieval. Default: ['mmr']
search_type (str, optional): 🔍 Search method for parameter exploration. Default: 'grid'. Options: 'grid', 'random', 'bayesian'
trials (int, optional): 🧪 Number of optimization trials. Default: 5
metric (str, optional): 📈 Metric to optimize. Default: 'faithfulness'
validation_choice (str, optional): ✅ Source of validation data. Default: 'generate'. Options: blank (use default), 'generate', local path, HF dataset ID
llm_model (str, optional): 🤖 LLM used for QA generation if validation_choice='generate'. Default: 'gemini-2.5-flash-lite'
"""
docs_path: Optional[str] = Field(
default="data/docs",
description="📂 Folder containing your documents for RAG optimization. Example: 'data/docs'"
)
retriever: Optional[List[str]] = Field(
description="🔍 Retriever type to use. Example: 'bm25', 'faiss', 'chroma'",
default=['faiss']
)
embedding_model: Optional[List[str]] = Field(
description="🧠 Embedding model name or path. Example: 'sentence-transformers/all-MiniLM-L6-v2'",
default=['sentence-transformers/all-MiniLM-L6-v2']
)
strategy: Optional[List[str]] = Field(
description="🎯 RAG strategy name. Example: 'fixed', 'token', 'sentence'",
default=['fixed']
)
chunk_sizes: Optional[List[int]] = Field(
description="📏 List of chunk sizes to evaluate. Example: [200, 400, 600]",
default=[200, 400, 600]
)
overlaps: Optional[List[int]] = Field(
description="🔁 List of overlap values to test. Example: [50, 100, 200]",
default = [50, 100, 200]
)
rerankers: Optional[List[str]] = Field(
default=["mmr"],
description="⚖️ Rerankers to apply after retrieval. Default: ['mmr']"
)
search_type: Optional[str] = Field(
default="grid",
description="🔍 Search method to explore parameter space. Options: 'grid', 'random', 'bayesian'"
)
trials: Optional[int] = Field(
default=5,
description="🧪 Number of optimization trials to run."
)
metric: Optional[str] = Field(
default="faithfulness",
description="📈 Evaluation metric for optimization. Options: 'faithfulness'"
)
validation_choice: Optional[str] = Field(
default='generate',
description=(
"✅ Validation data source. Options:\n"
" - Leave blank → use default 'validation_qa.json' if available\n"
" - 'generate' → auto-generate a validation QA file from your docs\n"
" - Path to a local JSON file (e.g. 'data/validation_qa.json')\n"
" - Hugging Face dataset ID (e.g. 'squad')"
)
)
llm_model: Optional[str] = Field(
default="gemini-2.5-flash-lite",
description="🤖 LLM used to generate QA dataset when validation_choice='generate'. Example: 'gemini-pro', 'gpt-4o-mini'"
)
class AutotuneRequest(BaseModel):
"""
⚡ Automatically tunes RAG pipeline parameters based on document analysis.
Args:
docs_path (str, optional): 📂 Folder containing documents for RAG optimization. Default: "data/docs"
embedding_model (str, optional): 🧠 Embedding model to analyze. Default: 'sentence-transformers/all-MiniLM-L6-v2'
num_chunk_pairs (int, optional): 🔢 Number of chunk pairs to analyze. Default: 5
metric (str, optional): 📈 Metric to optimize. Default: 'faithfulness'
search_type (str, optional): 🔍 Search method for parameter exploration. Default: 'grid'. Options: 'grid', 'random', 'bayesian'
trials (int, optional): 🧪 Number of optimization trials. Default: 5
validation_choice (str, optional): ✅ Source of validation data. Default: 'generate'. Options: blank, 'generate', local path, HF dataset ID
llm_model (str, optional): 🤖 LLM used for QA generation if validation_choice='generate'. Default: 'gemini-2.5-flash-lite'
"""
docs_path: Optional[str] = Field(
default="data/docs",
description="📂 Folder containing your documents for RAG optimization. Example: 'data/docs'"
)
embedding_model: Optional[str] = Field(
default="sentence-transformers/all-MiniLM-L6-v2",
description="🧠 Embedding model name or path. Example: 'sentence-transformers/all-MiniLM-L6-v2'"
)
num_chunk_pairs: Optional[int] = Field(
default=5,
description="🔢 Number of chunk pairs to analyze for tuning."
)
metric: Optional[str] = Field(
default="faithfulness",
description="📈 Evaluation metric for optimization. Options: 'faithfulness'"
)
search_type: Optional[str] = Field(
default="grid",
description="🔍 Search method to explore parameter space. Options: 'grid', 'random', 'bayesian'"
)
trials: Optional[int] = Field(
default=5,
description="🧪 Number of optimization trials to run."
)
validation_choice: Optional[str] = Field(
default='generate',
description=(
"✅ Validation data source. Options:\n"
" - Leave blank → use default 'validation_qa.jsonl' if available\n"
" - 'generate' → auto-generate a validation QA file from your docs\n"
" - Path to a local JSON file (e.g. 'data/validation_qa.json')\n"
" - Hugging Face dataset ID (e.g. 'squad')"
)
)
llm_model: Optional[str] = Field(
default="gemini-2.5-flash-lite",
description="🤖 LLM used to generate QA dataset when validation_choice='generate'. Example: 'gemini-pro', 'gpt-4o-mini'"
)
class QARequest(BaseModel):
"""
🧩 Generate a validation QA dataset from documents for RAG evaluation.
Args:
docs_path (str): 📂 Folder containing documents. Default: 'data/docs'
llm_model (str): 🤖 LLM model used for question generation. Default: 'gemini-2.5-flash-lite'
batch_size (int): 📦 Number of documents per batch. Default: 5
min_q (int): ❓ Minimum number of questions per document. Default: 3
max_q (int): ❓ Maximum number of questions per document. Default: 25
"""
docs_path: str = Field(
description="📂 Folder containing your documents to generate QA pairs from. Example: 'data/docs'",
default='data/docs'
)
llm_model: str = Field(
default="gemini-2.5-flash-lite",
description="🤖 LLM model used for question generation. Example: 'gemini-2.5-flash-lite', 'gpt-4o-mini'"
)
batch_size: int = Field(
default=5,
description="📦 Number of documents processed per generation batch."
)
min_q: int = Field(
default=3,
description="❓ Minimum number of questions per document."
)
max_q: int = Field(
default=25,
description="❓ Maximum number of questions per document."
)
|