<a href="https://colab.research.google.com/github/argonne-lcf/ai-science-training-series/blob/main/05_llm_part2/RAGTutorial.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# Retrieval Augmented Generation (RAG)

## Overview
*   Motivation for RAG
*   Idea behind RAG
*   Advantages and Disadvantages
*   Implementation to augment question + answer
*   Advanced applications

#### Imagine you went to live under a rock on August 2006. When you come out in 2024, you are asked how many planets revolve around the sun. What would you say?...
![pluto](https://github.com/architvasan/LLMWorkshop/blob/main/rag_images/pluto_planets.jpeg?raw=1)

This is similar to LLMs which are trained with data until a certain point and then asked questions on data they are not trained on. Understandably, LLMs will either be unable to answer or simply hallucinate a probably wrong answer.

###What can be done?

Have the LLM go to the library using **Retrieval Augmented Generation (RAG)**!

RAG involves adding your own data (via a retrieval tool) to the prompt that you pass into a large language model.

![rag architecture](https://github.com/architvasan/LLMWorkshop/blob/main/rag_images/rag-overview.original.png?raw=1)
Image credit: https://scriv.ai/guides/retrieval-augmented-generation-overview/

RAG has been shown to improve LLM prediction accuracy without needing to increase parameter size.

![rag architecture](https://github.com/architvasan/LLMWorkshop/blob/main/rag_images/rag_acc_v_size.png?raw=1)

*Image credit: Yu, Wenhao. "Retrieval-augmented generation across heterogeneous knowledge." Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies: Student Research Workshop. 2022.*

RAG also increases explainability by giving the source for information.

![rag architecture](https://github.com/architvasan/LLMWorkshop/blob/main/rag_images/rag_source_locator.png?raw=1)

Image credit: https://ai.stanford.edu/blog/retrieval-based-NLP/

## Advantages and Disadvantages

### Advantages

*   Provides domain specific context
*   Improves predictive performance and reduces hallucinations
*   Does not increase model parameters
*   Less labor intensive than fine-tuning LLMs

### Disadvantages

*   May introduce latency since we are adding a relatively costly search step
*   If your dataset includes private information, you may inadvertently expose another user with this information.
*   The data you want to use needs to be curated and you should decide how the data should be accessed. This adds time for the initial set-up.

#Implementation

### 1. Install + load relevant modules:
*   langchain
*   torch
*   transformers
*   sentence-transformers
*   datasets
*   faiss-cpu  
*   pypdf
*  unstructure[pdf]
*  huggingface_hub (add hf_token)



```{python}
#| colab: {base_uri: https://localhost:8080/}
#| scrolled: true
!pip install langchain==0.1.5
!pip install --quiet langchain_experimental
!pip install torch
!pip install transformers
!pip install faiss-cpu
!pip install pypdf
!pip install sentence-transformers
!pip install unstructured==0.12.3
!pip install unstructured[pdf]==0.12.3
!pip install tiktoken
!pip install huggingface_hub
from huggingface_hub import login

hf_token = "hf_yourtoken"
login(token=hf_token, add_to_git_credential=True)
```

```{python}
#| scrolled: true
#| colab: {base_uri: https://localhost:8080/}
!pip install -U langchain-community
```

### 2. Choose a dataset to use and then load it into your code
Here we are using the pdfs loaded in pdfs/. We load this using langchain DirectoryLoader.

We can load multiple types of datasets into this example though the most commonly used are PDFs and websites.

To load websites, we could also use `langchain WebBaseLoader`

In this example, we will consider PDFs and load them in using `langchain DirectoryLoader`.

We host all PDFs at the PDFs directory `ai-science-training-series/05_llm_part2/PDFs`


```{python}
#| colab: {base_uri: https://localhost:8080/}
! git clone https://github.com/argonne-lcf/ai-science-training-series.git
```

The below code is the langchain based document loader; however, recently this showed failure. We are working to fix it but we have an alternative code to run using bash commands

```{python}
#| scrolled: true
#| colab: {base_uri: https://localhost:8080/}
from langchain.document_loaders import DirectoryLoader
loader = DirectoryLoader('ai-science-training-series/05_llm_part2/PDFs', glob="**/*.pdf", show_progress=True)
documents = loader.load()
```

```{python}
# from pypdf import PdfReader
# x = !ls llm-workshop/tutorials/04-rag/PDFs/*.pdf
# print(x)
# # creating a pdf reader object
# documents = []
# for f in x:
#   reader = PdfReader('llm-workshop/tutorials/04-rag/PDFs/1703.07076.pdf')

#   # printing number of pages in pdf file
#   print(len(reader.pages))

#   # creating a page object
#   page = reader.pages[0]

#   # extracting text from page
#   # print(page.extract_text())
#   documents.append(page.extract_text())
```

### 3. Now, we need to split our documents into chunks.
We want the embedding to be greater than 1 word but much less than an entire page. This is essential for the similarity search between the query and the document. Essentially, the query will be searched for greatest similarity to embedded chunks in the dataset. Then those chunks with greatest similarity are augmented to the query.

It is essential to choose the chunking method according to your data type.
There are different ways to do this:

Fixed size
*   Token: Splits text on tokens. Can chunk tokens together
*   Character: Splits based on some user defined character.

Recursive
*  Recursively splits text. Useful for keeping related pieces of text next to each other.

Document based
*   HTML: Splits text based on HTML-specific characters.
*   Markdown: Splits on Markdown-specific characters
*   Code: Splits text based on characters specific to coding languages.

Semantic chunking
*   Extract semantic meaning from embeddings and then assess the semantic relationship between these chunks. Essentially splits into sentences, then groups into groups of 3 sentences, and then merges one that are similar in the embedding space.

Here we use recursive where the dataset is split using a set of characters. The default characters provided to it are ["\n\n", "\n", " ", ""].  A large text is split by the first character \n\n. If the first split by \n\n is still large then it moves to the next character which is \n and tries to split by it. This continues until the chunk size is reached.

```{python}
from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
#documents = text_splitter.create_documents(documents)
docs = text_splitter.split_documents(documents)
```

```{python}
#| scrolled: true
#| colab: {base_uri: https://localhost:8080/}
print(docs)
```


```{python}
#| scrolled: true
# from langchain.text_splitter import RecursiveCharacterTextSplitter
# text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
# docs = text_splitter.split_documents(documents)
```

### 4. Then we embed the chunked texts using a Transformer.
This allows us to encode the text into our search

Embedding converts text to a numerical representation in a vector space. RAG compares the embeddings of user queries within the vector of the knowledge library.

In this example, we choose a simple embedding using the MiniLM

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 458, referenced_widgets: [50ab3c6b796a417f82eae97fdff209a5, 92ead2409a454641895e0250bdcfef9e, 2736a3bcfeed4d0ca1e4b057b6ba263c, 96a48c244bde44f783e17d2fb09ab30a, 9d2dca32e0b04c7a830b2ea6b6709d3b, 1d0d1b4d8e664053b5ecc9ac18ae4d95, a351e18bdd13415a8bfc2925792144e8, 0a51867e85bb4ca091b342d3b444ba0e, 767b536a3bc546119278b81576d373f4, 944d4d22802d48a79e95bd454b227af0, 08d8aa99df3a4eca9beb3795df5a9896, 391271ffe92c4c5f9293704f0e44baa7, 4b8d5a71094e4961a9bb589dbc50fc91, 5fc35bc7eb0c477bac62296252729b14, 99e27e3724b149daaea6635184b47c4c, 1ecc8143b4d84c9a890d2f0b90e3ed01, f97a270a971f438483be1b55ad73108b, 683589aa5e8442e59a2fa2463cdb72d2, bcc4bdcaf218483e95e14789c121d065, a0dc6b70c3df41deaee6ffd21bac52c9, 5f8b307ac7954fe2a24179a409abb5fd, 0837819fda384c4cabd5276f62733c65, c98d05af6a244d6fa8e362fb5ff3cf4f, e59501eda6314551a58a62a376e2fb40, 0b03bd364ab74b63be48cd36638af1ad, 9481a37e853848fd95d65adffbdbdea6, 5df3f5f0058e48e1b9f0936989c10003, cdefd91bd9024e7980ce1b2e4de212a7, 116b5c29f916401b8cda95433769823b, 3eeb2e05033c41648e2465bb597c8d94, 5b3c747f8a7e4326b3eae01c5b809851, 04c5866455bc460c91672fdf5e62af49, 3a51a4709d5a455ab588e1121890726d, b956c7fa66e145bd9a6f1089b6fbd1e1, 5b328001cb684d2c9bfc91ed9f335141, 460195c92d83417191621463926bf878, bdf62004be78469cb077718421936ea7, 6b671e2331d54ef1a19e1622352ca8f0, 94f738ea2ecb4d589c7bc188a974f875, e01a0891183f41b9b8cb7df660ee6dcb, 209b5810fe114d6d8509265cf1aa42bb, e8d09f7ae86e4082856f5244b27962a9, 6b57d16116a04fa0a7e4769c79e19480, b0877b9497724e00bf4e539751a93cfe, 9516be91ee184931b721f27b5421fc39, d689173163b8414dabae7c8e070189c8, e1cd19fc345a46808ed2137c66e9f648, 114f8d5b692249a9a5e231409f42c347, 97df4be860e441ff82656478d73152d5, a3d56be0e1ed4157939479c14eb25d59, b307d2019b6548158dfd9eb2e2d81266, f807016b612e4c01aa0b4698b370a722, efa5edbd2d754094a1979c2b74f5d133, 6f57b73719b848c792e22ebed6301905, 41fc56daebe14c5baa8b36107d7175b8, b5c579dcc4c84dfb83aaabc416c6a9e7, 22fcd9377d244129acd96f46e3bc17e7, 5febe76d866741bda6fded73a471049b, b4119fda1ac14ca596654075c37ff420, 1c539a577435446ebf2d599e1f777630, 0b5751aeb9934365ae60003114e5f593, e8a92a347fd942bd8776b6167571bcf1, c7d7690e79c34e3b86fb898cecf2ca4c, 9584d33540a7442d937513363c72a0ab, 801dd358909843f0bbbada7d10099ca7, f03aacd78c5040619e5d89a9475af1d5, ab86031acdc749e9b31bc1698a635ce0, e2d4d8658aa94155bd4a92ad114c51b8, bed8e2726d8840d2b08a5f129dc19bc1, b9a70ae5d92a465b980845e87a67f1e4, c38415d2e45d4ff7ba871d99f9b1b523, fa81521a40d540abbdd2564fe70a954d, daca27b9821d4badb7243649d9a3c360, ecf056fc7070415fb83ea4deb3887583, 1fb2b2c4b2954480955dc530cc0e7905, 403be7357b3f40d6a8a6f428aed102f2, 57afef0e704a48e5b57752f5092070a8, 51e3d6c4048348f9bd6d878360774fea, 2d132b6c6841423080b69c7ea91b4c05, 7166a3209cf544308dd965da2939380a, 35d99e80f961460588173e44abf8cef8, b1a6e12f57b74c3892be626f1d088b49, 6accc604666140efad6a0b7b6054d3e3, 6b1cdd7ee89840a99f99c0bb865e58d2, 8062ce6a0fc1414e919c576e68dea7ce, de8aba8f805b4c0fb5579f23e9a9de5c, e7687eb21dec4e2cb049fd12a7a2b9a2, cc12a4910eec435ba623c63822953be4, ef2508cbbff842c28ca964a949c0f64a, 53478ed6b90c4e2195cd263f63923478, 1b03b0db59df438db1d73f0e893ab8da, b5146d3bbdc14737bb71f62dcd71ea24, dcbd16e6845a4eb68acf1f344bc2ff8f, 11fdd17b1f0e4a349e58f757586bd6c1, b4a5837db14048f28a54c9ec825f3440, d126a7834be44b6a81094298433aa126, 671f20dd28b2405f8fc8e3380d47f757, 3bb4bacbb11a4f40af7414cd10d3aec6, 2be6a69e298846f6ba91cfb3633eab1c, 6195f1f406a343e1b731efce4a5a95b1, 89e343f8574a4496931c22756f4bca83, 6e55bddd29af4579b721c16a2fc63e1c, 3441ed26851742819018b4abbe19e212, 807bca63da8e445099de076866aff172, 8ee3f6e207e04250b5320ff30088ab23, be489b0db48a403e8cbf868c372c2e4e, 29b6641a6c9448498b444d14e668221a, b412f6f12a84414eb0dc5e900671d436, f2509b738f9f4d94928304c5a46d86f4, 1df2254fe3504ed3a86c593869b0f9db, dcc44a136fbd412bac5fab9d4b06a40a, 508eba817d7d41b78808fbc040427596, 562538ee70f0404ca5c63c167e6d3a21, 4f7b9a64a59744d581f3bc26765643ff, 34285e769efe4e2e8095e15c739caf8e, 401b60be3839406987d29c995b3679da, 13fb95ab92de41b88347c7aed2e9441c, 85c7e3cb187448dca099364db4b805be, a72fa2953d4c4f70bcd75ec88c2ea95d, 1b737b62f51e4b1095b7d3db1c060efe, d31607e4830a4e0788624b459f809e96]}
from langchain.embeddings import HuggingFaceEmbeddings
modelPath = "sentence-transformers/all-MiniLM-l6-v2"
model_kwargs = {'device':'cpu'}
encode_kwargs = {'normalize_embeddings':False}
embeddings = HuggingFaceEmbeddings(
  model_name = modelPath,
  model_kwargs = model_kwargs,
  encode_kwargs=encode_kwargs
)
```

### 5.Create a vector database
Vector databases, also called vector storage, efficiently store and retrieve vector data, which are arrays of numerical values representing points in multi-dimensional space. They're useful for handling data like embeddings from deep learning models or numerical features. Unlike traditional relational databases, which aren't optimized for vectors, vector databases offer efficient storage, indexing, and querying for high-dimensional and variable-length vectors.

There are various types of vector databases:
1. Chroma
2. FAISS
3. Pinecone
4. Weaviate
5. Qdrant

Here, we build this using the FAISS utility.

![vector_database](https://github.com/argonne-lcf/llm-workshop/blob/main/tutorials/04-rag/rag_images/vector_database.png?raw=1)

Image credit: https://blog.gopenai.com/primer-on-vector-databases-and-retrieval-augmented-generation-rag-using-langchain-pinecone-37a27fb10546

```{python}
#| colab: {base_uri: https://localhost:8080/}
from langchain.vectorstores import FAISS
db = FAISS.from_documents(docs, embeddings)
question = "What is RF Fold?"
searchDocs = db.similarity_search(question)
print(searchDocs[0].page_content)
```

### 6. Initialize the LLM that will be used for question answering

Here, we use a pretrained model flan-t5-large as part of a HuggingFacePipeline. This will later be chained with the vector database for RAG.

```{python}
#| scrolled: true
#| colab: {base_uri: https://localhost:8080/, height: 313, referenced_widgets: [c4804d4a46df4aa1b9b089e48e54aa82, 9cbdb36c80f546c0b9e066c0e3ad85c1, 7b22e2c90d174f5dad664e5995c425b1, 88457ec03982476eb958efb1ea2d5085, a4a99d57c30b44699f70a5b7dfe366a8, 609c953b06e14252ad4999f3e9c6aa2b, 5f37d315b0744d20b8b35e61eeda8e2c, e87a73d92853477f9ac1bd767c11999a, ee6c4bf481b04b3284116fc2849a14d5, f8c0492545f54a3aac8b2bf6f68f146b, 715d28a3fb5f429dbe21a05513584aef, db5a066686154f05997080e98ce59abc, b23437e99a5a4e75830d4762ac89da9c, 3f5317ca9c0f4f09894326c9532c71cf, 95948e44862a441a819712cd71ca6eae, 89066b2795b44e50905e6e0544e094b0, f1f08e9c066f46ee8e2352e144e26f0a, 2ca48ed368fc417f9483facb6a3a58ae, d199fb112d4c4417a85efef894288068, 488cca4d9fb7456186b3fdc8a9e5a51a, 37cc1f7ee94f4424a4d4bca31caa09ec, 6e6182ce57bf4c61a87d800556392705, 90ba1c5997c54789b3e306db6bf7cdb6, 9836b480662a47919f33ee38ac5e6236, 94b701eceb2f4ee7a1c3eacd75676ae5, 5748cc4f24444c989993535513a427fb, e0fd0716509740928f1543ec07b08b12, c7978b7c1fca42a4a67b2532ebebbcd0, b417dff54be9488fa6f94bc06be183f4, 7dd24d7223b243bbbb56dcd7ccd78c7a, 70afe58926c245f9be76995560405afc, d45a45e2036f42d1984e047ebd63fa32, 57b0c1b425e34c3caa2d0cd8521a2911, db16259e7cba44dfa636304774c2c6db, cf2c8f83a6b34cd88f7cebd0175c72af, 847851e1f2214e1cb3d66f1cc2d3405e, aa13b0a1cc8a4b709f4dd9479a472ed1, 966a23bf2cfd4619822d48be8c4701c6, a2bec6964c5f4320929dd0ae77e56a41, 41e85494f1ec496d8fcd359017decbd0, bf1bac252f6945df860a53417d1942d6, 73d92a1009b542fd8d919846785f380b, 9b5d4a5edbf640d58d554fb7ce8acc58, 04cc1b7ac3a34690b3851f47efb8815e, 68bd97c79f7b4bceb203828a758d107e, 800a7753b75744bdb02173766a825515, 6e5e135b276e4a4583f016851faf5d08, 96016c6eb931439aaa01a1a76e7bd376, ec15c2c66d0e4c5e8fe32dc9ddfd5034, 08db420bd45e4f508676d784df456fbc, 1e51258337e148e7ae9a408d4286f2b9, 4edf8bda68ce45d9aaa266e83314e0f8, fbd3f1b8b2104b56a792ac72504c40d2, 4fb225f5989c402fbb88bc95946fae47, b7ab29f24a6743ee8e7dec6f6ad5725f, 4e5131daf3fb4645a35c098f3db3f74b, 35d4ac78182f4dbc8cd28f886998e2ba, a3e334873908452fa418b7d39b933c20, 32d2ad1b1205484694970c7becefa6dd, 62f90487c5d94e7fbeca07bbfb77aca2, 6174d061c9a840928eb8e5a64d3764d8, eb23c29ff7c74084ab1566389561a9cd, a4f24f1fbe15430bb93621c45ef2cc95, 25c12ded834343ebbd53247a3bce686b, 9dfefb4431ff4989871b19e88e28cf33, 78e12ba3a9ab463ebca1bce5019d1edc, ae96aaef5684415f818413c0b3f2a4a5, f7ec65d9bb134a8eaec3e5476d406e13, 25c1bf7af8aa4c9886edcdb45ca1fe4d, da29190b171641d78390ff070fe40701, 250e2c59d8cd4b12b74a03d84e6370af, 86470b106c59468e96056b638b2ce696, 9ac1876586274f9789ab2530d499aff5, 7c8aeae0703c43ea9b2c60dd07700099, 366d99b698f743b485ba40bafe7575e8, 1cce0cdc28e9442e8d334408e22c14d7, 0c1aaadc4d514c878f6be620ebb0d6e7]}
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,pipeline
from langchain import HuggingFacePipeline

tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
llm = HuggingFacePipeline(
   pipeline = pipe,
   model_kwargs={"temperature": 0, "max_length": 2048},
)
#'HuggingFaceH4/zephyr-7b-beta'
```

### 7. Retrieve data and use it to answer a question

![rag_workflow](https://github.com/argonne-lcf/llm-workshop/blob/main/tutorials/04-rag/rag_images/rag_workflow.png?raw=1)

Image credit: https://blog.gopenai.com/retrieval-augmented-generation-101-de05e5dc21ef

Let's ask questions it would only be able to know if the model actually read the texts!

```{python}
from langchain.prompts import PromptTemplate

template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Keep the answer as concise as possible.
{context}
Question: {question}
Helpful Answer:"""
QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
```

```{python}
#| colab: {base_uri: https://localhost:8080/}
from langchain.chains import RetrievalQA
qa_chain = RetrievalQA.from_chain_type(
  llm=llm,
  chain_type="stuff",
  retriever=db.as_retriever(),
  chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
)
result = qa_chain ({ "query" : "What technique proposed in 2023 can be used to predict protein folding?" })
print(result["result"])
```

Now let's ask the chain where to find the article related to RFDiffusion

```{python}
#| colab: {base_uri: https://localhost:8080/}
qa_chain ({ "query" : "Which scientific article should I read to learn about RFdiffusion for protein folding?" })
```

## Exercise

Use any of the frameworks/models here to load in your favorite websites and ask the model a question regarding them.

Hint:
```
from langchain_community.document_loaders import WebBaseLoader

loader = WebBaseLoader(["https://www.espn.com/", "https://www.google.com"])

To bypass SSL verification errors during fetching, you can set the “verify” option:

loader.requests_kwargs = {‘verify’:False}

data = loader.load()
```


## More applications

### RAG using Llama 2, Langchain and ChromaDB

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 1000}
#| scrolled: true
!pip install einops xformers \
bitsandbytes chromadb
```

```{python}
#| colab: {base_uri: https://localhost:8080/}
#| scrolled: true
!pip install accelerate
```

Load in models and setup pipeline

```{python}
from torch import cuda, bfloat16
import torch
import transformers
from transformers import AutoTokenizer
from time import time
#import chromadb
#from chromadb.config import Settings
from langchain.llms import HuggingFacePipeline
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.vectorstores import Chroma

model_id = 'meta-llama/Llama-2-7b-chat-hf'

device = f'cuda'

# set quantization configuration to load large model with less GPU memory
# this requires the `bitsandbytes` library
bnb_config = transformers.BitsAndBytesConfig(
    load_in_4bit=False,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=bfloat16
)
```

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 368, referenced_widgets: [8df15ee07d5b4f319135edbd1336f6bc, b23c31e839cb4eb2912511925fd89b9f, 351778cfee7548368c3efd2a9fa45c4d, 40b0a738218b4d0d9b0f0136b61f5475, c3e6a516af6b413aa722862306800aec, 34e7d064f786428b90ef133c026c1996, d4620904edf142a492f2d7952336161b, 29cb8b36ac7a4c29bb03635d3a327673, 7cb5dc7be34a4c1f978ae0276eb056e1, 0e6e6b9ea3e0413aa348c204209be5de, db9629950d8e4c779195afc8e1fdd121, c4255efba6014f098e33cfcf5c2b0baf, e5234b26e1634d239033398d900856a6, bbb80dc24edd474db1ad8d52a11aee55, c806af04c8bd42f29d722bab4b446290, bfa0651a56b447c3b273615630382196, dadb95c3acfc4c2899ceccb6d4b27864, b3d019eb94554f62827a6e971597a15c, c13aec6e4d0e44bfa32acc18bdf3f203, d50455156a5a4d45a73a50cb7f416531, f62635f7b37c43dda0f4a30794e60ed2, 8c60c4ae54cd476cb5b916dd9be59020, 24f381d25e1246fcb1dda99772a86a6d, 52897dd8be92400799b2e8faafff0280, 6da5c9f6a2c3416699a8a2d42a16700b, 093543555c124ec3b83d67d81ec0a894, e3dc4f7b5f434308b8a3979f367e5cb5, c83fdc64f38447c8899d0d63e4838f5b, 7623d841ae044f4b8acc549e8a703c82, 6cd26551e9904703bfad3d2b285bd995, 68326900690d4293826fe4df969488b7, a15a50db09f849389a40645d3632aa30, 588b2fe072634cf88afb8ffb7f5c193f, 7042a066086143d5b581cbeab2e5a28f, 82c936a81cd241138b620e485619fb24, f6d86fda3aec497aa7acf4451db0ed54, 56da550bc1794a9b98a4cce8c6e10c29, 1e51efd76d014283afbbe40da2e6551f, 6d2ef60932cc4eb98cad2c597b4b4530, 2559f7a616504ecebb0044ec1011e99a, a1d4ef4d59bc49f0941422717e9ebccc, e8ada17a8a8d40a5b60b4921897907aa, a9bd2ae265154b769be5b5a23d9aa23e, 89d1eae4a1df4ec3b778328f7686caba, 5f10b36d4d4d41ad94d6bf9da79c1156, ec28be3ab6324055b4ecbdb0912bbe68, 34a1ceeaf38345a4a1893fc6e9c05db7, 309a455e823445828819289f656210ee, 648f07aa15784e0b9d87f773f57cb4d2, 7bd68f3b0bab4f548b2aa6f927650b7b, 9a4f9bfb7bc64e83b66757c34875c756, 4365308bdcf04d6f8f369093c2d321ba, 9bec35b479314cbcaae2bbe11b9f2bdf, abadb02290ce45abbcf152949b3734a1, 176988f6991d4d71b5c9a33321d2153e, 055004f61d074c27bca1d7fc5acb793b, 98022d05b6754ac09eef55e0d99e24fc, fc4865f0d11d4d83bf4871cdde6cd556, 34bffb01bbf947f6b4aa0fa1c74474ae, d8d5cb60c8f2420e8ffed156dcb9ccff, 0f72aca74fa74f8eb9241636c390cb06, b6d1b668476048fabab9faeb90a79fd2, 15675159e5614d3aaeb1ac9a09b35599, e24c7a7154c74190954aa79c72424b4a, 84d0d9b2a32c4317936bcc508c19599d, 542b4ce411244f0581eb3c164807ab03]}
#| scrolled: true
time_1 = time()
model_config = transformers.AutoConfig.from_pretrained(
    model_id,
)
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    config=model_config,
    quantization_config=bnb_config,
    device_map='auto',
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
time_2 = time()
print(f"Prepare model, tokenizer: {round(time_2-time_1, 3)} sec.")
```

```{python}
#| colab: {base_uri: https://localhost:8080/}
time_1 = time()
query_pipeline = transformers.pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        torch_dtype=torch.float16,
        device_map="auto",)
time_2 = time()
print(f"Prepare pipeline: {round(time_2-time_1, 3)} sec.")
```

Prepare function to test pipeline

```{python}
def test_model(tokenizer, pipeline, prompt_to_test):
    """
    Perform a query
    print the result
    Args:
        tokenizer: the tokenizer
        pipeline: the pipeline
        prompt_to_test: the prompt
    Returns
        None
    """
    # adapted from https://huggingface.co/blog/llama2#using-transformers
    time_1 = time()
    sequences = pipeline(
        prompt_to_test,
        do_sample=True,
        top_k=10,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        max_length=200,)
    time_2 = time()
    print(f"Test inference: {round(time_2-time_1, 3)} sec.")
    for seq in sequences:
        print(f"Result: {seq['generated_text']}")
```

```{python}
#| colab: {base_uri: https://localhost:8080/}
test_model(tokenizer,
           query_pipeline,
           "Please explain what is the State of the Union address. Give just a definition. Keep it in 100 words.")
```

Set up Huggingface pipeline

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 211}
llm = HuggingFacePipeline(pipeline=query_pipeline)
# checking again that everything is working fine
llm(prompt="Please explain what is the State of the Union address. Give just a definition. Keep it in 100 words.")
```

Load in data

```{python}
loader = TextLoader("llm-workshop/tutorials/04-rag/state_union2023.txt",
                    encoding="utf8")
documents = loader.load()
```

Chunk data recursively

```{python}
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=20)
all_splits = text_splitter.split_documents(documents)
```

```{python}
#| colab: {base_uri: https://localhost:8080/}
#| scrolled: true
all_splits
```

Embed and store in Chroma Vector store

```{python}
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device':'cpu'}

embeddings = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs)
```

```{python}
vectordb = Chroma.from_documents(documents=all_splits, embedding=embeddings, persist_directory="chroma_db")
```

Set up chain

```{python}
retriever = vectordb.as_retriever()

qa = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    verbose=True
)
```

Test RAG

```{python}
def test_rag(qa, query):
    print(f"Query: {query}\n")
    time_1 = time()
    result = qa.run(query)
    time_2 = time()
    print(f"Inference time: {round(time_2-time_1, 3)} sec.")
    print("\nResult: ", result)
```

Get sources...

```{python}
docs = vectordb.similarity_search(query)
print(f"Query: {query}")
print(f"Retrieved documents: {len(docs)}")
for doc in docs:
    doc_details = doc.to_json()['kwargs']
    print("Source: ", doc_details['metadata']['source'])
    print("Text: ", doc_details['page_content'], "\n")
```

```{python}
query = "What were the main topics in the State of the Union in 2023? Summarize. Keep it under 200 words."
test_rag(qa, query)
```

