Views
No views yet
pip install htmlragpip install -e .1from htmlrag import clean_html
2
3question = "When was the bellagio in las vegas built?"
4html = """
5<html>
6<head>
7<h1>Bellagio Hotel in Las</h1>
8</head>
9<body>
10<p class="class0">The Bellagio is a luxury hotel and casino located on the Las Vegas Strip in Paradise, Nevada. It was built in 1998.</p>
11</body>
12<div>
13<div>
14<p>Some other text</p>
15<p>Some other text</p>
16</div>
17</div>
18<p class="class1"></p>
19<!-- Some comment -->
20<script type="text/javascript">
21document.write("Hello World!");
22</script>
23</html>
24"""
25
26#. alternatively you can read html files and merge them
27# html_files=["/path/to/html/file1.html", "/path/to/html/file2.html"]
28# htmls=[open(file).read() for file in html_files]
29# html = "\n".join(htmls)
30
31simplified_html = clean_html(html)
32print(simplified_html)
33
34# <html>
35# <h1>Bellagio Hotel in Las</h1>
36# <p>The Bellagio is a luxury hotel and casino located on the Las Vegas Strip in Paradise, Nevada. It was built in 1998.</p>
37# <div>
38# <p>Some other text</p>
39# <p>Some other text</p>
40# </div>
41# </html>1# Maximum number of words in a node when constructing the block tree for pruning with the embedding model
2MAX_NODE_WORDS_EMBED = 10
3# MAX_NODE_WORDS_EMBED = 256 # a recommended setting for real-world HTML documents
4# Maximum number of tokens in the output HTML document pruned with the embedding model
5MAX_CONTEXT_WINDOW_EMBED = 60
6# MAX_CONTEXT_WINDOW_EMBED = 6144 # a recommended setting for real-world HTML documents
7# Maximum number of words in a node when constructing the block tree for pruning with the generative model
8MAX_NODE_WORDS_GEN = 5
9# MAX_NODE_WORDS_GEN = 128 # a recommended setting for real-world HTML documents
10# Maximum number of tokens in the output HTML document pruned with the generative model
11MAX_CONTEXT_WINDOW_GEN = 32
12# MAX_CONTEXT_WINDOW_GEN = 4096 # a recommended setting for real-world HTML documents1from htmlrag import build_block_tree
2
3block_tree, simplified_html = build_block_tree(simplified_html, max_node_words=MAX_NODE_WORDS_EMBED)
4# block_tree, simplified_html = build_block_tree(simplified_html, max_node_words=MAX_NODE_WORDS_GEN, zh_char=True) # for Chinese text
5for block in block_tree:
6 print("Block Content: ", block[0])
7 print("Block Path: ", block[1])
8 print("Is Leaf: ", block[2])
9 print("")
10
11# Block Content: <h1>Bellagio Hotel in Las</h1>
12# Block Path: ['html', 'title']
13# Is Leaf: True
14#
15# Block Content: <div>
16# <p>Some other text</p>
17# <p>Some other text</p>
18# </div>
19# Block Path: ['html', 'div']
20# Is Leaf: True
21#
22# Block Content: <p>The Bellagio is a luxury hotel and casino located on the Las Vegas Strip in Paradise, Nevada. It was built in 1998.</p>
23# Block Path: ['html', 'p']
24# Is Leaf: True1from htmlrag import EmbedHTMLPruner
2
3embed_model = "BAAI/bge-large-en"
4query_instruction_for_retrieval = "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: "
5embed_html_pruner = EmbedHTMLPruner(embed_model=embed_model, local_inference=True, query_instruction_for_retrieval = query_instruction_for_retrieval)
6# alternatively you can init a remote TEI model, refer to https://github.com/huggingface/text-embeddings-inference.
7# tei_endpoint="http://YOUR_TEI_ENDPOINT"
8# embed_html_pruner = EmbedHTMLPruner(embed_model=embed_model, local_inference=False, query_instruction_for_retrieval = query_instruction_for_retrieval, endpoint=tei_endpoint)
9block_rankings=embed_html_pruner.calculate_block_rankings(question, simplified_html, block_tree)
10print(block_rankings)
11
12# [2, 0, 1]
13
14#. alternatively you can use bm25 to rank the blocks
15from htmlrag import BM25HTMLPruner
16bm25_html_pruner = BM25HTMLPruner()
17block_rankings = bm25_html_pruner.calculate_block_rankings(question, simplified_html, block_tree)
18print(block_rankings)
19
20# [2, 0, 1]
21
22from transformers import AutoTokenizer
23
24chat_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-70B-Instruct")
25
26pruned_html = embed_html_pruner.prune_HTML(simplified_html, block_tree, block_rankings, chat_tokenizer, MAX_CONTEXT_WINDOW_EMBED)
27print(pruned_html)
28
29# <html>
30# <h1>Bellagio Hotel in Las</h1>
31# <p>The Bellagio is a luxury hotel and casino located on the Las Vegas Strip in Paradise, Nevada. It was built in 1998.</p>
32# </html>1from htmlrag import GenHTMLPruner
2import torch
3
4# construct a finer block tree
5block_tree, pruned_html = build_block_tree(pruned_html, max_node_words=MAX_NODE_WORDS_GEN)
6# block_tree, pruned_html = build_block_tree(pruned_html, max_node_words=MAX_NODE_WORDS_GEN, zh_char=True) # for Chinese text
7for block in block_tree:
8 print("Block Content: ", block[0])
9 print("Block Path: ", block[1])
10 print("Is Leaf: ", block[2])
11 print("")
12
13# Block Content: <h1>Bellagio Hotel in Las</h1>
14# Block Path: ['html', 'title']
15# Is Leaf: True
16#
17# Block Content: <p>The Bellagio is a luxury hotel and casino located on the Las Vegas Strip in Paradise, Nevada. It was built in 1998.</p>
18# Block Path: ['html', 'p']
19# Is Leaf: True
20
21ckpt_path = "zstanjj/HTML-Pruner-Llama-1B"
22if torch.cuda.is_available():
23 device="cuda"
24else:
25 device="cpu"
26gen_embed_pruner = GenHTMLPruner(gen_model=ckpt_path, device=device)
27block_rankings = gen_embed_pruner.calculate_block_rankings(question, pruned_html, block_tree)
28print(block_rankings)
29
30# [1, 0]
31
32pruned_html = gen_embed_pruner.prune_HTML(pruned_html, block_tree, block_rankings, chat_tokenizer, MAX_CONTEXT_WINDOW_GEN)
33print(pruned_html)
34
35# <p>The Bellagio is a luxury hotel and casino located on the Las Vegas Strip in Paradise, Nevada. It was built in 1998.</p>| Dataset | ASQA | HotpotQA | NQ | TriviaQA | MuSiQue | ELI5 |
|---|---|---|---|---|---|---|
| Metrics | EM | EM | EM | EM | EM | ROUGE-L |
| BM25 | 49.50 | 38.25 | 47.00 | 88.00 | 9.50 | 16.15 |
| BGE | 68.00 | 41.75 | 59.50 | 93.00 | 12.50 | 16.20 |
| E5-Mistral | 63.00 | 36.75 | 59.50 | 90.75 | 11.00 | 16.17 |
| LongLLMLingua | 62.50 | 45.00 | 56.75 | 92.50 | 10.25 | 15.84 |
| JinaAI Reader | 55.25 | 34.25 | 48.25 | 90.00 | 9.25 | 16.06 |
| HtmlRAG-Phi-3.8B | 68.50 | 46.25 | 60.50 | 93.50 | 13.25 | 16.33 |
| HtmlRAG-Llama-1B | 66.50 | 45.00 | 60.75 | 93.00 | 10.00 | 16.25 |
1@misc{tan2024htmlraghtmlbetterplain,
2 title={HtmlRAG: HTML is Better Than Plain Text for Modeling Retrieved Knowledge in RAG Systems},
3 author={Jiejun Tan and Zhicheng Dou and Wen Wang and Mang Wang and Weipeng Chen and Ji-Rong Wen},
4 year={2024},
5 eprint={2411.02959},
6 archivePrefix={arXiv},
7 primaryClass={cs.IR},
8 url={https://arxiv.org/abs/2411.02959},
9}