ajuste design pattern

rag_utils.py

@@ -17,32 +17,46 @@ device = f'cuda:{torch.cuda.current_device()}' if torch.cuda.is_available() else
 
 import os
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
-cache_dir = "/home/user/.cache/huggingface" #"./model/qwen-awq" #"/home/felipe/.cache/huggingface/transformers" #"/home/user/.cache/huggingface"
+# cache_dir = "/home/user/.cache/huggingface" #"./model/qwen-awq" #"/home/felipe/.cache/huggingface/transformers" #"/home/user/.cache/huggingface"
 
-embedding_model = HuggingFaceEmbeddings(model_name=config.local_emb_path)
+_embedding_instance = None
+_model_instance = None
+_tokenizer = None
+
+def get_embedding_model():
+    global _embedding_instance
+    if _embedding_instance is None:
+        if config.local_emb_path is None:
+            raise ValueError("⚠️ config.local_emb_path ainda não foi inicializado!")
+        _embedding_instance = HuggingFaceEmbeddings(model_name=config.local_emb_path)
+    return _embedding_instance
 
 # model_name = "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8" #"Qwen/Qwen2.5-7B-Instruct-AWQ" #"Qwen/Qwen2.5-7B-Instruct"
-model = AutoModelForCausalLM.from_pretrained(
-    config.local_model_path,
-    torch_dtype="auto",
-    device_map="auto",
-    trust_remote_code=True,
-    # cache_dir=cache_dir
-)
-model.to(device)
-tokenizer = AutoTokenizer.from_pretrained(config.local_model_path, trust_remote_code=True)#, cache_dir=cache_dir)
-
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    temperature=0.1,
-    do_sample=False
-)
 
-# Adapta para LangChain
-llm = HuggingFacePipeline(pipeline=pipe)
+def get_model():
+    global _model_instance
+    if _model_instance is None:
+        if config.local_model_path is None:
+            raise ValueError("⚠️ config.local_model_path ainda não foi inicializado!")
+        _model_instance = AutoModelForCausalLM.from_pretrained(
+            config.local_model_path,
+            torch_dtype="auto",
+            device_map="auto",
+            trust_remote_code=True
+        )
+    
+    return _model_instance
+
+# _model_instance.to(device)
+
+def get_tokenizer():
+    global _tokenizer
+    if _tokenizer is None:
+        if config.local_model_path is None:
+            raise ValueError("⚠️ config.local_model_path ainda não foi inicializado!")
+        _tokenizer = AutoTokenizer.from_pretrained(config.local_model_path, trust_remote_code=True)
+    
+    return _tokenizer
 
 def create_split_doc(raw_text):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
@@ -51,6 +65,7 @@ def create_split_doc(raw_text):
     return docs
 
 def store_docs(docs):
+    embedding_model = get_embedding_model()
     vectorstore = FAISS.from_documents(docs, embedding_model)
     return vectorstore
 
@@ -73,14 +88,26 @@ Pergunta:
     return prompt_template
 
 def create_rag_chain(vectorstore):
+    pipe = pipeline(
+        "text-generation",
+        model=get_model(),
+        tokenizer=get_tokenizer(),
+        max_new_tokens=512,
+        temperature=0.1,
+        do_sample=False
+    )
+
+    # Adapta para LangChain
+    llm = HuggingFacePipeline(pipeline=pipe)
+
     rag_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    retriever=vectorstore.as_retriever(),
-    chain_type="stuff",
-    chain_type_kwargs={"prompt": create_template()}
-)
+        llm=llm,
+        retriever=vectorstore.as_retriever(),
+        chain_type="stuff",
+        chain_type_kwargs={"prompt": create_template()}
+    )
+    
     return rag_chain
 
-
 if __name__ == '__main__':
     pass