enhanced performance and readibility

Author: VinciGit00

examples/local_models/pdf_scraper_ollama.py

@@ -5,7 +5,7 @@
 
 graph_config = {
     "llm": {
-        "model": "ollama/llama3",
+        "model": "ollama/mistral",
         "temperature": 0,
         "format": "json",  # Ollama needs the format to be specified explicitly
         "model_tokens": 4000,

scrapegraphai/nodes/generate_answer_node.py

@@ -99,7 +99,8 @@ def execute(self, state: dict) -> dict:
                     input_variables=["question"],
                     partial_variables={"context": chunk.page_content,
                                        "format_instructions": format_instructions})
-                
+                chain =  prompt | self.llm_model | output_parser
+                answer = chain.invoke({"question": user_prompt})
             else:
                 prompt = PromptTemplate(
                     template=template_chunks,
@@ -125,11 +126,7 @@ def execute(self, state: dict) -> dict:
             )
             merge_chain = merge_prompt | self.llm_model | output_parser
             answer = merge_chain.invoke({"context": answer, "question": user_prompt})
-        else:
-            # Chain
-            single_chain = list(chains_dict.values())[0]
-            answer = single_chain.invoke({"question": user_prompt})
-
         # Update the state with the generated answer
         state.update({self.output[0]: answer})
-        return state
+        return state
+       

scrapegraphai/nodes/generate_answer_pdf_node.py

@@ -117,6 +117,9 @@ def execute(self, state):
                         "format_instructions": format_instructions,
                     },
                 )
+
+                chain =  prompt | self.llm_model | output_parser
+                answer = chain.invoke({"question": user_prompt})
             else:
                 prompt = PromptTemplate(
                     template=template_chunks_pdf,
@@ -145,10 +148,6 @@ def execute(self, state):
             )
             merge_chain = merge_prompt | self.llm_model | output_parser
             answer = merge_chain.invoke({"context": answer, "question": user_prompt})
-        else:
-            # Chain
-            single_chain = list(chains_dict.values())[0]
-            answer = single_chain.invoke({"question": user_prompt})
 
         # Update the state with the generated answer
         state.update({self.output[0]: answer})