File size: 2,450 Bytes
982eab4 d4c3edd d665d4a d4c3edd e2765f4 d665d4a e2765f4 d665d4a d4c3edd e2765f4 d4c3edd 982eab4 d665d4a d4c3edd d665d4a d4c3edd 982eab4 d665d4a d4c3edd d665d4a 982eab4 d4c3edd d665d4a d4c3edd d665d4a d4c3edd 982eab4 d665d4a d4c3edd d665d4a d4c3edd d665d4a d4c3edd d665d4a 982eab4 d665d4a d4c3edd d665d4a d4c3edd d665d4a d4c3edd d665d4a 982eab4 e2765f4 d4c3edd |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
from typing import List
from langchain.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.docstore.document import Document
from langchain.chains import RetrievalQA
from langchain_community.llms import HuggingFaceHub
from langchain_huggingface import HuggingFaceEmbeddings
# Configure safe cache directories (writable within container)
CACHE_DIR = "/tmp/huggingface"
os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
os.environ["HF_HOME"] = CACHE_DIR
os.makedirs(CACHE_DIR, exist_ok=True)
# Constants
DATA_PATH = "/app/data"
VECTORSTORE_PATH = "/app/vectorstore"
DOCS_FILENAME = "context.txt"
VECTORSTORE_INDEX_NAME = "faiss_index"
EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-MiniLM-L6-v2"
LLM_REPO_ID = "mistralai/Mistral-7B-Instruct-v0.1"
def load_embedding_model() -> Embeddings:
"""Load Hugging Face sentence transformer embeddings."""
return HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
def load_documents() -> List[Document]:
"""Load documents and split them into manageable chunks."""
loader = TextLoader(os.path.join(DATA_PATH, DOCS_FILENAME))
documents = loader.load()
splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
return splitter.split_documents(documents)
def load_vectorstore() -> FAISS:
"""Load FAISS vectorstore from disk or create it from documents."""
vectorstore_dir = os.path.join(VECTORSTORE_PATH, VECTORSTORE_INDEX_NAME)
embedding_model = load_embedding_model()
if os.path.exists(vectorstore_dir):
return FAISS.load_local(
folder_path=vectorstore_dir,
embeddings=embedding_model,
allow_dangerous_deserialization=True,
)
documents = load_documents()
vectorstore = FAISS.from_documents(documents, embedding_model)
vectorstore.save_local(vectorstore_dir)
return vectorstore
def ask_question(query: str) -> str:
"""Run a question-answering chain with the retriever and language model."""
vectorstore = load_vectorstore()
retriever = vectorstore.as_retriever()
llm = HuggingFaceHub(
repo_id=LLM_REPO_ID,
model_kwargs={"temperature": 0.5, "max_tokens": 256},
)
qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
return qa_chain.run(query)
|