File size: 1,988 Bytes
ac8c947 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
from langchain.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import TextLoader
from langchain_huggingface import HuggingFacePipeline
from langchain.chains import RetrievalQA
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
def load_documents(file_path: str):
"""Loads documents from a specified file path."""
loader = TextLoader(file_path)
return loader.load()
def split_documents(documents, chunk_size=500, chunk_overlap=50):
"""Splits documents into chunks."""
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
return splitter.split_documents(documents)
def create_embeddings(model_name="sentence-transformers/all-MiniLM-L6-v2"):
"""Creates HuggingFace embeddings."""
return HuggingFaceEmbeddings(model_name=model_name)
def setup_vector_store(docs, embeddings, persist_directory="./chroma_db"):
"""Sets up and persists the Chroma vector store."""
db = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
return db.as_retriever()
def create_qa_chain(retriever, model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0"):
"""Creates the RetrievalQA chain with streaming capabilities."""
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
)
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=512,
temperature=0.7,
top_p=0.9
)
llm = HuggingFacePipeline(pipeline=pipe)
qa_chain = RetrievalQA.from_chain_type(
llm=llm,
retriever=retriever,
chain_type="stuff",
return_source_documents=True # Added to potentially help with streaming or understanding context
)
return qa_chain
|