Dataset Generation

from rag_opt.rag import DatasetGenerator
from rag_opt import init_chat_model

# Initialize LLM for generation
llm = init_chat_model(
    model="gpt-3.5-turbo",
    model_provider="openai",
    api_key="sk-***"
)

# Create generator with your documents
generator = DatasetGenerator(
    llm=llm,
    dataset_path="./data"  # Path to your documents
)

# Generate dataset
dataset = generator.generate(
    n=6,  # Number of QA pairs
    batch_size=2
)

# Save for later use
dataset.to_json("./rag_dataset.json")

from langchain.schema import Document
from rag_opt import init_chat_model
from rag_opt.rag import Parser

# Use custom parser
parser = Parser(path="./data")
docs = parser.load_docs()

# Or provide documents directly
documents = [
    Document(page_content="Your text here..."),
    Document(page_content="More content...")
]

llm = init_chat_model(
    model="gpt-3.5-turbo",
    model_provider="openai",
    api_key="sk-***"
)

generator = DatasetGenerator(
    llm=llm,
    parser=parser  # or pass documents to generate()
)

dataset = generator.generate(
    n=5,
    source_docs=documents
)

# NOTE:: u have to provide only contexts and difficulty_instruction in single bracket
custom_prompt = """
Generate question-answer pairs from these contexts:
{contexts}

Difficulty: {difficulty_instruction}

Output as JSON array: [{{"question": "...", "answer": "..."}}]
"""

generator = DatasetGenerator(
    llm=llm,
    dataset_path="./data",
    custom_prompt=custom_prompt
)

import asyncio

dataset = await generator.agenerate(
    n=10,
    batch_size=2
)

{
  "items": [
    {
      "question": "What is the capital of France?",
      "answer": "Paris",
      "contexts": ["France is a country...", "Paris is located..."],
    },
    ...
  ]
}

from rag_opt.dataset import TrainDataset

# Load from JSON
dataset = TrainDataset.from_json("./rag_dataset.json")

# Access items
for item in dataset.items:
    print(f"Q: {item.question}")
    print(f"A: {item.answer}")