samim2024 commited on
Commit
d4c3edd
·
verified ·
1 Parent(s): d665d4a

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +32 -24
model.py CHANGED
@@ -1,62 +1,70 @@
1
  import os
 
2
  from langchain.vectorstores import FAISS
3
- from langchain_huggingface import HuggingFaceEmbeddings
4
  from langchain_community.document_loaders import TextLoader
5
  from langchain.text_splitter import CharacterTextSplitter
6
  from langchain.docstore.document import Document
7
  from langchain.chains import RetrievalQA
8
  from langchain_community.llms import HuggingFaceHub
9
- from langchain.embeddings.base import Embeddings
10
 
11
- # Set safe caching directories to avoid permission denied errors
12
- os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
13
- os.environ["HF_HOME"] = "/app/cache"
14
- os.makedirs("/app/cache", exist_ok=True)
 
15
 
16
  # Constants
17
  DATA_PATH = "/app/data"
18
  VECTORSTORE_PATH = "/app/vectorstore"
19
  DOCS_FILENAME = "context.txt"
 
20
  EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-MiniLM-L6-v2"
 
21
 
22
 
23
  def load_embedding_model() -> Embeddings:
24
- """Initialize and return the HuggingFace embedding model."""
25
  return HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
26
 
27
 
28
- def load_documents() -> list[Document]:
29
- """Load and split documents into chunks."""
30
  loader = TextLoader(os.path.join(DATA_PATH, DOCS_FILENAME))
31
- raw_docs = loader.load()
 
32
  splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
33
- docs = splitter.split_documents(raw_docs)
34
- return docs
35
 
36
 
37
  def load_vectorstore() -> FAISS:
38
- """Load or create FAISS vectorstore from documents."""
39
- vectorstore_file = os.path.join(VECTORSTORE_PATH, "faiss_index")
40
-
41
  embedding_model = load_embedding_model()
42
 
43
- if os.path.exists(vectorstore_file):
44
- return FAISS.load_local(vectorstore_file, embedding_model, allow_dangerous_deserialization=True)
 
 
 
 
45
 
46
- docs = load_documents()
47
- vectorstore = FAISS.from_documents(docs, embedding_model)
48
- vectorstore.save_local(vectorstore_file)
49
  return vectorstore
50
 
51
 
52
  def ask_question(query: str) -> str:
53
- """Query the vectorstore and return the answer using the language model."""
54
  vectorstore = load_vectorstore()
 
55
 
56
  llm = HuggingFaceHub(
57
- repo_id="mistralai/Mistral-7B-Instruct-v0.1",
58
  model_kwargs={"temperature": 0.5, "max_tokens": 256},
59
  )
60
 
61
- qa = RetrievalQA.from_chain_type(llm=llm, retriever=vectorstore.as_retriever())
62
- return qa.run(query)
 
1
  import os
2
+ from typing import List
3
  from langchain.vectorstores import FAISS
4
+ from langchain.embeddings.base import Embeddings
5
  from langchain_community.document_loaders import TextLoader
6
  from langchain.text_splitter import CharacterTextSplitter
7
  from langchain.docstore.document import Document
8
  from langchain.chains import RetrievalQA
9
  from langchain_community.llms import HuggingFaceHub
10
+ from langchain_huggingface import HuggingFaceEmbeddings
11
 
12
+ # Configure safe cache directories (writable within container)
13
+ CACHE_DIR = "/tmp/huggingface"
14
+ os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
15
+ os.environ["HF_HOME"] = CACHE_DIR
16
+ os.makedirs(CACHE_DIR, exist_ok=True)
17
 
18
  # Constants
19
  DATA_PATH = "/app/data"
20
  VECTORSTORE_PATH = "/app/vectorstore"
21
  DOCS_FILENAME = "context.txt"
22
+ VECTORSTORE_INDEX_NAME = "faiss_index"
23
  EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-MiniLM-L6-v2"
24
+ LLM_REPO_ID = "mistralai/Mistral-7B-Instruct-v0.1"
25
 
26
 
27
  def load_embedding_model() -> Embeddings:
28
+ """Load Hugging Face sentence transformer embeddings."""
29
  return HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
30
 
31
 
32
+ def load_documents() -> List[Document]:
33
+ """Load documents and split them into manageable chunks."""
34
  loader = TextLoader(os.path.join(DATA_PATH, DOCS_FILENAME))
35
+ documents = loader.load()
36
+
37
  splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
38
+ return splitter.split_documents(documents)
 
39
 
40
 
41
  def load_vectorstore() -> FAISS:
42
+ """Load FAISS vectorstore from disk or create it from documents."""
43
+ vectorstore_dir = os.path.join(VECTORSTORE_PATH, VECTORSTORE_INDEX_NAME)
 
44
  embedding_model = load_embedding_model()
45
 
46
+ if os.path.exists(vectorstore_dir):
47
+ return FAISS.load_local(
48
+ folder_path=vectorstore_dir,
49
+ embeddings=embedding_model,
50
+ allow_dangerous_deserialization=True,
51
+ )
52
 
53
+ documents = load_documents()
54
+ vectorstore = FAISS.from_documents(documents, embedding_model)
55
+ vectorstore.save_local(vectorstore_dir)
56
  return vectorstore
57
 
58
 
59
  def ask_question(query: str) -> str:
60
+ """Run a question-answering chain with the retriever and language model."""
61
  vectorstore = load_vectorstore()
62
+ retriever = vectorstore.as_retriever()
63
 
64
  llm = HuggingFaceHub(
65
+ repo_id=LLM_REPO_ID,
66
  model_kwargs={"temperature": 0.5, "max_tokens": 256},
67
  )
68
 
69
+ qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
70
+ return qa_chain.run(query)