07 · Lab: chat with your documents
Open in Colab View the notebook on GitHub
The outputs on this page are the ones committed with the notebook. Without an API key, the notebook replays recorded model answers, and the same recordings appear here.
Session 01 · Building AI applications for research and data · the lab, after notebook 6
Objective
Build a small application that answers questions from documents of your choice, and use it. You start with the eight documents of the session, so that everything runs at once, and then you can replace them with your own: reports, notes, papers or a data dictionary. The documents stay inside the notebook. Nothing is sent to an outside service, and no key is needed.
Every piece comes from a notebook you have already seen:
| Piece | From |
|---|---|
| Split the documents into chunks and store them | Notebook 02 |
| Retrieve the passages closest to a question, and say so when none is close | Notebook 03 |
| Ask a model to answer from those passages | Notebooks 01 and 03 |
| Look up every number of the answer in the passages, and mark it | The idea of notebook 06, on documents |
The application is deliberately simple: one question at a time, with no memory of the conversation. The enhancements that a full chatbot adds are listed at the end of the notebook.
The model is new: a small open one, Qwen2.5-1.5B-Instruct, that runs in the notebook itself. It is far weaker than the models behind the chatbots you know. That is useful here, because its mistakes show why the checks of the earlier notebooks are needed.
The application
Contents
- Section 1 chooses the documents: the session's eight, or your own.
- Section 2 splits the documents into chunks and stores them.
- Section 3 loads the model that runs inside the notebook.
- Section 4 asks one question from code: retrieve, answer, and check every number.
- Section 5 asks three questions that show what the checks can and cannot tell you.
- Section 6 opens the chat box, where you ask your own questions.
How to run it
This notebook runs a real language model inside Colab, so it needs more than the others. No account key and no payment are needed. Do these steps in this order:
- Choose a GPU first. In the Colab menu, click Runtime > Change runtime type, pick T4 GPU, and click Save. Do this before you run anything, because changing it later restarts the notebook. A GPU makes each answer take a few seconds instead of up to a minute. If Colab says no GPU is available, the notebook still works on the CPU, only more slowly.
- Run everything. Click Runtime > Run all. If Colab warns that the notebook comes from GitHub, click Run anyway.
- Wait for the first run. It downloads the model, about 3 GB, and takes a few minutes. The cell of step 3 prints "loaded on the GPU" when it is ready. Cells that show a spinner are still working.
- Chat. Scroll to step 6 at the end, type a question in the box, and press Enter or click Ask. The first answer is the slowest.
If something goes wrong, click Runtime > Restart session, then run all again. The chat box in step 6 appears only in a running Colab or Jupyter notebook, so it is empty when you read the saved copy. To use your own documents, set USE_MY_DOCUMENTS = True in step 1 before you run all, and upload your files when asked.
# Setup. In Colab this clones the repository and installs the dependencies; locally it uses the folder as is.
import os, sys, subprocess, pathlib
REPO_URL = "https://github.com/worldbank/ai4data.git" # the program's repository
BRANCH = "main" # the branch the Colab setup clones
SESSION = "training/session-01-building-ai-applications/notebooks"
if not pathlib.Path("workshop.py").exists():
subprocess.run(["git", "clone", "-q", "--depth", "1", "--branch", BRANCH, REPO_URL, "/content/ai4data"], check=True)
os.chdir(f"/content/ai4data/{SESSION}")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-r", "requirements.txt"], check=True)
sys.path.insert(0, os.getcwd())
import workshop as w
print("Ready.")
Ready.
1. Choose the documents
The session documents are used unless you say otherwise. The cell lists them. To use your own, set USE_MY_DOCUMENTS = True and run the cell again: Colab asks you to pick files. Plain text, Markdown and PDF files work, and a zip of several files is unpacked. A scanned PDF has no text to read, so it will not work. If you would like to see what a document looks like first, lab.download_session_documents() saves the eight session documents as a zip file, which you can open, change and upload again.
import lab
USE_MY_DOCUMENTS = False # True: upload your own files (Colab asks you to pick them)
paths = lab.upload_documents() if USE_MY_DOCUMENTS else lab.session_documents()
w.show_table(["document", "size"], [(p.name, f"{p.stat().st_size / 1000:.1f} KB") for p in paths], source=f"{len(paths)} documents")
| document | size |
|---|---|
| coffee.md | 1.1 KB |
| internet-indicator.md | 2.1 KB |
| kenya-digital.md | 1.7 KB |
| kenya-economy.md | 1.2 KB |
| mobile-money.md | 1.5 KB |
| regional-2016.md | 1.1 KB |
| rwanda-digital.md | 1.4 KB |
| rwanda-economy.md | 1.2 KB |
8 documents
2. Split and store
As in notebook 02, each document is split into chunks, by section when it has headings and in windows of about 700 characters when it does not, and every chunk is stored with its meaning, so that a question can find it. The table shows what was made.
chunks = lab.prepare(paths)
store = w.build_store(chunks)
lab.chunk_table(chunks)
| document | chunks | characters per chunk |
|---|---|---|
| coffee.md | 1 | 888 |
| internet-indicator.md | 4 | 414 |
| kenya-digital.md | 4 | 325 |
| kenya-economy.md | 3 | 285 |
| mobile-money.md | 1 | 1,159 |
| regional-2016.md | 1 | 730 |
| rwanda-digital.md | 4 | 265 |
| rwanda-economy.md | 3 | 268 |
21 chunks in all
3. The model
The cell downloads the model on the first run and loads it. If Colab has given you a GPU, the model goes onto it and the cell says which one; if not, it runs on the CPU and the cell tells you how to switch. Because the model runs here, the documents are never sent anywhere. This cell has no saved output, because the message depends on the machine that runs it.
llm = lab.load_model()
chat = lab.DocumentChat(llm, store, k=4, min_score=0.25)
Qwen/Qwen2.5-1.5B-Instruct loaded on the CPU, so answers will take up to a minute. For quick answers: Runtime > Change runtime type > T4 GPU, then run the notebook again from the top.
4. One question, asked from code
chat.ask(question) does the whole round: it retrieves the closest passages, drops any that are not close enough (min_score), asks the model to answer from what is left, and looks up every number of the answer in those passages. The answer is shown as a chat. A ✓ after a number means that the number appears in one of the passages; hover over it to see where. A ⚠ means that it does not. Years and citations are not checked. Under the answer, the passages that the answer was written from can be opened.
from IPython.display import HTML, display
display(HTML(lab.reply_html(chat.ask("What share of Rwanda's population used the internet in 2024?"))))
The 4 passages the answer was written from
# Rwanda: internet use and mobile subscriptions, 2010 to 2024 ## Internet use The share of individuals using the internet in Rwanda was 8.0 percent of the population in 2010 and 12.5 percent in 2015. It reached 23.8 percent in 2020, 28.0 percent in 2022, 28.8 percent in 2023, and 31.7 percent in 2024, the latest year available. The period of fastest growth was 2015 to 2020, when the share nearly doubled. Growth…
# Regional digital summary, 2016 Internet use in East Africa continues to expand. In Kenya, 16.6 percent of the population used the internet in 2015, up from 7.2 percent in 2010. In Rwanda the share was 12.5 percent in 2015, up from 8.0 percent in 2010. Mobile cellular subscriptions reached 80.1 per 100 people in Kenya and 75.3 per 100 people in Rwanda in 2015. On current trends, internet use in both countries …
## Population Rwanda's population was 10.3 million in 2010 and 14.3 million in 2024.
## Reading the two indicators together Rwanda recorded growth above 7 percent in 10 of the 15 years of the period; 2020 was the only year of contraction. GDP per capita in current dollars rose by about 80 percent between 2010 and 2024 from a lower starting level than Kenya's.
5. What the checks can and cannot tell you
Three questions on the session documents show the three situations that you will meet. They are asked from code, so that the answers are saved in this copy of the notebook.
for question in [
"What is the capital of France?",
"How have Kenya's coffee export earnings changed in recent years?",
"What is Kenya's GDP per capita and how does it compare with Rwanda's?",
]:
display(HTML(lab.reply_html(chat.ask(question))))
The 4 passages the answer was written from
# Coffee exports in East Africa Coffee is one of the traditional export crops of Kenya and Rwanda. Both countries grow mainly arabica varieties at altitude, and both have moved part of their production towards specialty grades that sell at a premium in export markets. In Kenya, coffee is grown by cooperatives of smallholders and by estates, and most of the crop is sold through a central auction. In Rwanda, washin…
## Reading the two indicators together Kenya's economy grew in every year of the period except 2020. GDP per capita in current dollars roughly doubled between 2010 and 2024, with the exchange rate accounting for part of the year-to-year variation.
# Kenya: GDP per capita and growth, 2010 to 2024 ## GDP per capita GDP per capita in current US dollars was 1,092 in 2010, 1,489 in 2015, and 1,928 in 2020. It reached 2,110 in 2022, fell to 1,943 in 2023, and stood at 2,133 in 2024. Because the series is in current dollars, movements reflect exchange rates and prices as well as real output; the 2023 decline coincides with a depreciation of the shilling.
# Kenya: internet use and mobile subscriptions, 2010 to 2024 ## Internet use The share of individuals using the internet in Kenya rose from 7.2 percent of the population in 2010 to 16.6 percent in 2015 and 24.7 percent in 2020. Growth continued after 2020: 29.4 percent in 2022, 32.1 percent in 2023, and 35.0 percent in 2024, the latest year available. Over the full period the share increased by about 28 percent…
The 4 passages the answer was written from
## Reading the two indicators together Rwanda recorded growth above 7 percent in 10 of the 15 years of the period; 2020 was the only year of contraction. GDP per capita in current dollars rose by about 80 percent between 2010 and 2024 from a lower starting level than Kenya's.
# Rwanda: GDP per capita and growth, 2010 to 2024 ## GDP per capita GDP per capita in current US dollars was 587 in 2010, 726 in 2015, and 803 in 2020. It passed 1,000 in 2022 at 1,007, reached 1,061 in 2023, and was 1,060 in 2024. The series is in current dollars, so exchange rates and prices affect it as well as real output.
# Kenya: GDP per capita and growth, 2010 to 2024 ## GDP per capita GDP per capita in current US dollars was 1,092 in 2010, 1,489 in 2015, and 1,928 in 2020. It reached 2,110 in 2022, fell to 1,943 in 2023, and stood at 2,133 in 2024. Because the series is in current dollars, movements reflect exchange rates and prices as well as real output; the 2023 decline coincides with a depreciation of the shilling.
## Reading the two indicators together Kenya's economy grew in every year of the period except 2020. GDP per capita in current dollars roughly doubled between 2010 and 2024, with the exchange rate accounting for part of the year-to-year variation.
- No close passage. The documents say nothing about France, so no passage reaches
min_scoreand the model is not asked. The application says so instead of guessing. This is the "no relevant result" case of notebook 03, handled by the application and not by the model. - A close passage with no numbers. The coffee note is close to the question, so the model answers. It has no figures to give, and a model that is weak may say more than the note does. No mark can appear, because there are no numbers, and nothing here checks the sentences. This is what a check on numbers does not cover.
- Numbers. The marks say whether each number is in a passage. They do not say that it is the right number for the question. In the saved copy of the third answer, the model wrote that Rwanda's GDP per capita "had risen to $803" by 2024. The figure 803 is in the Rwanda document, so it carries a ✓, but it is the figure for 2020: the value for 2024 is 1,060. The check found the number and could not see that the sentence puts it in the wrong year. Your own run may word the answers differently, and a small model makes this kind of slip often. In notebook 06 each number named its own record, with its country and year, which is the stronger check; here the check is what can be done on arbitrary documents.
6. Chat
Now use the application. Type a question in the box under the cell and press Ask (or Enter). Some to try on the session documents: What is mobile money? How did Kenya's GDP per capita change between 2010 and 2020? Which sources mention coffee? With your own documents, ask what you would ask a colleague who had read them. The chat box appears when you run the cell, so it is empty in a saved copy of this notebook.
Each question is answered on its own. The application does not remember the earlier ones, so a follow-up such as "and for Rwanda?" has nothing to refer to; write the full question each time.
lab.launch(chat)
A complete version: the Data360 chat
What you built here is the small version. The World Bank's Data360 chat is a complete version that puts the ideas of this series to work in one application: a language model that answers questions about development data, tools that fetch the records, and numbers that can be checked against them (the protocol of notebook 06). Open it, ask the main question of the session, and compare what you see with your own chat box: where does each answer come from, and how does it show that a number can be trusted?
What this application leaves out
This is the simplest version that works, and it leaves out what a full chatbot has. Each of these is an enhancement that you can add later, and none is needed for a useful first version:
- Memory of the conversation, so that follow-up questions work.
- Managing the context: the application sends the four closest passages and the question, and nothing more. It does not count how much text the model can take, and it does not shorten or summarise long passages.
- Speed-ups: reusing the model's work on text it has already read (a KV cache), and showing the answer as it is written.
- Better search: re-ranking the passages, combining a search by words with the search by meaning, or rewriting the question before it is searched.
- Keeping the index, so that the documents are not split and stored again each time the notebook starts.
- A test suite of your own questions, as in section 6 of notebook 03 for retrieval, to see whether a change made the application better.
Key points
- An application that chats with documents is the pieces of the earlier notebooks put together: split and store, retrieve, answer from the passages, check.
- A small model runs on your own machine, and the documents stay there. It is also weaker, so the checks and the retrieval threshold carry more of the work.
- Saying "I could not find that" when no passage is close is the application's job, and costs nothing.
- A check on numbers finds numbers that are not in the sources. It cannot find a claim that has no number.
- The Data360 chat is the complete version of this application, with the enhancements below.
Exercise
- Set
USE_MY_DOCUMENTS = True, upload two or three files of your own, and rerun from step 1. Ask one question that the files answer and one that they do not. What does the application do with the second? - After step 3, run
chat.min_score = 0.5, then ask the coffee question again in the chat box of step 6. What changes, and what is the risk of a high threshold? Try0.1as well. - Run
chat.k = 1, then paste the main question of the session (w.MAIN_QUESTION) into the chat box. How much of the answer is left, and why? - On a GPU runtime, rerun from step 3 with
lab.load_model("Qwen/Qwen2.5-3B-Instruct"). Do the answers get better? Which of the marks change?
Where this connects to the program
- Practical Guide to Serving Official Statistics to AI Agents: the design of tools over a statistical catalog, in full.
- Practical Guide to Evaluation Suites for Statistical AI: how to build and run the tests.
- Proof-Carrying Numbers: the claim-tag protocol the verifier in this session follows. See also the paper and the
pcnrepository. - Data360 MCP server: the tools of notebook 04, exposed to any MCP client over the World Bank's data platform.
- End-to-end reference implementation: documentation, validation, publication, agent access, and verification in one run.