feat: adicionando embedding personalizados

compose.yaml

@@ -7,6 +7,7 @@ services:
       - SECRET_KEY=${SECRET_KEY}
       - DATABASE_PASSWORD=${DATABASE_PASSWORD}
       - OPENAI_API_KEY=${OPENAI_API_KEY}
+      - HUGGINGFACEHUB_API_TOKEN=${HUGGINGFACEHUB_API_TOKEN}
     env_file:
       - .env
     develop:

endpoint_teste/serializer.py

@@ -17,4 +17,5 @@ class PDFUploadSerializer(serializers.Serializer):
     file = serializers.FileField()
     system_prompt = serializers.CharField(required=True)
     user_message = serializers.CharField(required=True)
-    model = serializers.CharField(required=False)
+    model = serializers.CharField(required=False)
+    embedding = serializers.CharField(required=False)

endpoint_teste/views.py

@@ -68,6 +68,9 @@ def getPDF(request):
             print('data: ', data)
             pdf_file = serializer.validated_data['file']
             pdf_file.seek(0)
+            
+            embedding = serializer.validated_data.get("embedding", "gpt")
+            model = serializer.validated_data.get("model", default_model)
             # print(dir(pdf_file))
             # print('pdf_file: ', pdf_file.read())
             # pdf_content = pdf_file.read()
@@ -87,10 +90,7 @@ def getPDF(request):
             print('temp_file_path: ', temp_file_path)
 
             resposta_llm = None
-            try:
-                resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path, model=serializer.validated_data['model'])
-            except:
-                resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path, model=default_model)
+            resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path, model=model, embedding=embedding)
 
             os.remove(temp_file_path)
 

langchain_backend/main.py

@@ -1,27 +1,38 @@
 import os
-from langchain_backend.utils import create_prompt_llm_chain, create_retriever, getPDF, vectorstore
+from langchain_backend.utils import create_prompt_llm_chain, create_retriever, getPDF
 from langchain_backend import utils
 from langchain.chains import create_retrieval_chain
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_chroma import Chroma
+from langchain_openai import OpenAIEmbeddings
 
 os.environ.get("OPENAI_API_KEY")
 
-def get_llm_answer(system_prompt, user_prompt, pdf_url, model):
+def get_llm_answer(system_prompt, user_prompt, pdf_url, model, embedding):
+  if embedding == "gpt":
+    embedding_object = OpenAIEmbeddings()
+  else:
+    embedding_object = HuggingFaceEmbeddings(model_name=embedding)
+
+  vectorstore = Chroma(
+      collection_name="documents",
+      embedding_function=embedding_object
+    )
+
+
   print('model: ', model)
+  print('embedding: ', embedding)
   pages = []
   if pdf_url:
     pages = getPDF(pdf_url)
   else:
     pages = getPDF()
-  retriever = create_retriever(pages)
-  # rag_chain = None
+  retriever = create_retriever(pages, vectorstore)
   rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt, model))
-  # if model:
-    # rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt, model))
-  # else:
-    # rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt))
   results = rag_chain.invoke({"input": user_prompt})
   print('allIds ARQUIVO MAIN: ',  utils.allIds)
   vectorstore.delete( utils.allIds)
+  vectorstore.delete_collection()
   utils.allIds = []
   print('utils.allIds: ', utils.allIds)
   return results

langchain_backend/utils.py

@@ -2,21 +2,18 @@ from langchain_community.document_loaders import PyPDFLoader
 import os
 from langchain_openai import ChatOpenAI
 from langchain_chroma import Chroma
-from langchain_openai import OpenAIEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
 from setup.environment import default_model
 from uuid import uuid4
 
+
 os.environ.get("OPENAI_API_KEY")
 os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
-vectorstore = Chroma(
-    collection_name="documents",
-    embedding_function=OpenAIEmbeddings()
-  )
 allIds = []
 
 def getPDF(file_path):
@@ -25,22 +22,16 @@ def getPDF(file_path):
   loader = PyPDFLoader(file_path, extract_images=False)
   pages = loader.load_and_split(text_splitter)
   for page in pages:
-    print('\n\n\n')
+    print('\n')
     print('allIds: ', allIds)
     documentId = str(uuid4())
     allIds.append(documentId)
     page.id = documentId
   return pages
 
-def create_retriever(documents):
+def create_retriever(documents, vectorstore):
   print('\n\n')
-  print('documents: ', documents)
-  
-#   vectorstore = Chroma.from_documents(
-#     documents,
-#     embedding=OpenAIEmbeddings(),
-# )
-  # vectorstore.delete_collection()
+  print('documents: ', documents[:2])
 
   vectorstore.add_documents(documents=documents)
 
@@ -58,12 +49,10 @@ def create_prompt_llm_chain(system_prompt, modelParam):
     model = HuggingFaceEndpoint(
         repo_id=modelParam,
         task="text-generation",
-        max_new_tokens=100,
+        # max_new_tokens=100,
         do_sample=False,
         huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
     )
-    # result = model.invoke("Hugging Face is")
-    # print('result: ', result)
 
   system_prompt = system_prompt + "\n\n" + "{context}"
   prompt = ChatPromptTemplate.from_messages(