@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 225 lines 6.71 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """RAG (Retrieval Augmented Generation) Implementation with LlamaIndex
4 ================================================================
5
6 This script demonstrates a RAG system using:
7 - LlamaIndex: For document indexing and retrieval
8 - Milvus: As vector store backend
9 - vLLM: For embedding and text generation
10
11 Features:
12 1. Document Loading & Processing
13 2. Embedding & Storage
14 3. Query Processing
15
16 Requirements:
17 1. Install dependencies:
18 pip install llama-index llama-index-readers-web \
19 llama-index-llms-openai-like \
20 llama-index-embeddings-openai-like \
21 llama-index-vector-stores-milvus \
22
23 2. Start services:
24 # Start embedding service (port 8000)
25 vllm serve ssmits/Qwen2-7B-Instruct-embed-base
26
27 # Start chat service (port 8001)
28 vllm serve qwen/Qwen1.5-0.5B-Chat --port 8001
29
30 Usage:
31 python retrieval_augmented_generation_with_llamaindex.py
32
33 Notes:
34 - Ensure both vLLM services are running before executing
35 - Default ports: 8000 (embedding), 8001 (chat)
36 - First run may take time to download models
37
38 """
39
40 import argparse
41 from argparse import Namespace
42 from typing import Any
43
44 from llama_index.core import Settings, StorageContext, VectorStoreIndex
45 from llama_index.core.node_parser import SentenceSplitter
46 from llama_index.embeddings.openai_like import OpenAILikeEmbedding
47 from llama_index.llms.openai_like import OpenAILike
48 from llama_index.readers.web import SimpleWebPageReader
49 from llama_index.vector_stores.milvus import MilvusVectorStore
50
51
52 def init_config(args: Namespace):
53 """Initialize configuration with command line arguments."""
54 return {
55 "url": args.url,
56 "embedding_model": args.embedding_model,
57 "chat_model": args.chat_model,
58 "vllm_api_key": args.vllm_api_key,
59 "embedding_endpoint": args.embedding_endpoint,
60 "chat_endpoint": args.chat_endpoint,
61 "db_path": args.db_path,
62 "chunk_size": args.chunk_size,
63 "chunk_overlap": args.chunk_overlap,
64 "top_k": args.top_k,
65 }
66
67
68 def load_documents(url: str) -> list:
69 """Load and process web documents."""
70 return SimpleWebPageReader(html_to_text=True).load_data([url])
71
72
73 def setup_models(config: dict[str, Any]):
74 """Configure embedding and chat models."""
75 Settings.embed_model = OpenAILikeEmbedding(
76 api_base=config["embedding_endpoint"],
77 api_key=config["vllm_api_key"],
78 model_name=config["embedding_model"],
79 )
80
81 Settings.llm = OpenAILike(
82 model=config["chat_model"],
83 api_key=config["vllm_api_key"],
84 api_base=config["chat_endpoint"],
85 context_window=128000,
86 is_chat_model=True,
87 is_function_calling_model=False,
88 )
89
90 Settings.transformations = [
91 SentenceSplitter(
92 chunk_size=config["chunk_size"],
93 chunk_overlap=config["chunk_overlap"],
94 )
95 ]
96
97
98 def setup_vector_store(db_path: str) -> MilvusVectorStore:
99 """Initialize vector store."""
100 sample_emb = Settings.embed_model.get_text_embedding("test")
101 print(f"Embedding dimension: {len(sample_emb)}")
102 return MilvusVectorStore(uri=db_path, dim=len(sample_emb), overwrite=True)
103
104
105 def create_index(documents: list, vector_store: MilvusVectorStore):
106 """Create document index."""
107 storage_context = StorageContext.from_defaults(vector_store=vector_store)
108 return VectorStoreIndex.from_documents(
109 documents,
110 storage_context=storage_context,
111 )
112
113
114 def query_document(index: VectorStoreIndex, question: str, top_k: int):
115 """Query document with given question."""
116 query_engine = index.as_query_engine(similarity_top_k=top_k)
117 return query_engine.query(question)
118
119
120 def get_parser() -> argparse.ArgumentParser:
121 """Parse command line arguments."""
122 parser = argparse.ArgumentParser(description="RAG with vLLM and LlamaIndex")
123
124 # Add command line arguments
125 parser.add_argument(
126 "--url",
127 default=("https://docs.vllm.ai/en/latest/getting_started/quickstart.html"),
128 help="URL of the document to process",
129 )
130 parser.add_argument(
131 "--embedding-model",
132 default="ssmits/Qwen2-7B-Instruct-embed-base",
133 help="Model name for embeddings",
134 )
135 parser.add_argument(
136 "--chat-model", default="qwen/Qwen1.5-0.5B-Chat", help="Model name for chat"
137 )
138 parser.add_argument(
139 "--vllm-api-key", default="EMPTY", help="API key for vLLM compatible services"
140 )
141 parser.add_argument(
142 "--embedding-endpoint",
143 default="http://localhost:8000/v1",
144 help="Base URL for embedding service",
145 )
146 parser.add_argument(
147 "--chat-endpoint",
148 default="http://localhost:8001/v1",
149 help="Base URL for chat service",
150 )
151 parser.add_argument(
152 "--db-path", default="./milvus_demo.db", help="Path to Milvus database"
153 )
154 parser.add_argument(
155 "-i", "--interactive", action="store_true", help="Enable interactive Q&A mode"
156 )
157 parser.add_argument(
158 "-c",
159 "--chunk-size",
160 type=int,
161 default=1000,
162 help="Chunk size for document splitting",
163 )
164 parser.add_argument(
165 "-o",
166 "--chunk-overlap",
167 type=int,
168 default=200,
169 help="Chunk overlap for document splitting",
170 )
171 parser.add_argument(
172 "-k", "--top-k", type=int, default=3, help="Number of top results to retrieve"
173 )
174
175 return parser
176
177
178 def main():
179 # Parse command line arguments
180 args = get_parser().parse_args()
181
182 # Initialize configuration
183 config = init_config(args)
184
185 # Load documents
186 documents = load_documents(config["url"])
187
188 # Setup models
189 setup_models(config)
190
191 # Setup vector store
192 vector_store = setup_vector_store(config["db_path"])
193
194 # Create index
195 index = create_index(documents, vector_store)
196
197 if args.interactive:
198 print("\nEntering interactive mode. Type 'quit' to exit.")
199 while True:
200 # Get user question
201 question = input("\nEnter your question: ")
202
203 # Check for exit command
204 if question.lower() in ["quit", "exit", "q"]:
205 print("Exiting interactive mode...")
206 break
207
208 # Get and print response
209 print("\n" + "-" * 50)
210 print("Response:\n")
211 response = query_document(index, question, config["top_k"])
212 print(response)
213 print("-" * 50)
214 else:
215 # Single query mode
216 question = "How to install vLLM?"
217 response = query_document(index, question, config["top_k"])
218 print("-" * 50)
219 print("Response:\n")
220 print(response)
221 print("-" * 50)
222
223
224 if __name__ == "__main__":
225 main()