|
from llama_index.core import ( |
|
VectorStoreIndex, |
|
SimpleDirectoryReader, |
|
get_response_synthesizer, |
|
ServiceContext, |
|
) |
|
from llama_index.embeddings.huggingface import HuggingFaceEmbedding |
|
from llama_index.core.postprocessor import SentenceTransformerRerank |
|
from typing import Optional, List |
|
from llama_index.llms.groq import Groq |
|
|
|
|
|
class RAG: |
|
def __init__( |
|
self, filepaths: List[str], rerank: Optional[SentenceTransformerRerank] = None |
|
) -> None: |
|
documents = SimpleDirectoryReader(input_files=filepaths).load_data() |
|
response_synthesizer = get_response_synthesizer( |
|
response_mode="tree_summarize", |
|
use_async=True, |
|
) |
|
self.index = VectorStoreIndex.from_documents( |
|
documents=documents, |
|
response_synthesizer=response_synthesizer, |
|
) |
|
if not rerank: |
|
self.query_engine = self.index.as_query_engine( |
|
response_mode="tree_summarize", |
|
use_async=True, |
|
streaming=True, |
|
similarity_top_k=10, |
|
) |
|
else: |
|
self.query_engine = self.index.as_query_engine( |
|
response_mode="tree_summarize", |
|
use_async=True, |
|
streaming=True, |
|
similarity_top_k=10, |
|
node_postprocessors=[rerank], |
|
) |
|
|
|
def run_query_engine(self, prompt): |
|
response = self.query_engine.query(prompt) |
|
response.print_response_stream() |
|
return str(response) |
|
|
|
|
|
class ServiceContextModule: |
|
def __init__(self, api_key, model_name) -> None: |
|
self._llm = Groq(model=model_name, api_key=api_key) |
|
self._embedding_model = HuggingFaceEmbedding( |
|
"Snowflake/snowflake-arctic-embed-m-long", trust_remote_code=True |
|
) |
|
self.service_context = ServiceContext.from_defaults( |
|
llm=self._llm, |
|
embed_model=self._embedding_model, |
|
) |
|
|