from llama_index.core.llama_dataset.generator import RagDatasetGenerator
from llama_index.llms.ollama import Ollama
from llama_index.core.schema import TextNode

node = TextNode(
    text="New York, often called New York City or simply NYC, "
    "is the most populous city in the United States, located at the "
    "southern tip of New York State on one of the world's largest "
    "natural harbors. The city comprises five boroughs, each of which "
    "is coextensive with a respective county. New York is a global center "
    "of finance and commerce, culture and technology, entertainment and "
    "media, academics and scientific output, and the arts and fashion, and, "
    "as home to the headquarters of the United Nations, is an important "
    "center for international diplomacy. New York City is the center of "
    "the world's principal metropolitan economy."
)

llm = Ollama(model="gemma:2b", request_timeout=60.0)
num_questions_per_chunk = 1
for i in range(10):
    gen = RagDatasetGenerator(
        nodes=[node],
        llm=llm,
        num_questions_per_chunk=num_questions_per_chunk,
    )
    generated_questions = gen.generate_questions_from_nodes()
    print(
        f"Number of questions expected: {num_questions_per_chunk} - "
        f"Number of questions generated: {len(generated_questions.examples)}"
    )