use llama3.1 for structuring query

openml-labs · Aug 9, 2024 · 06b3b77 · 06b3b77
1 parent 4f2a095
commit 06b3b77
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/structured_query/llm_service_structured_query.py b/structured_query/llm_service_structured_query.py
@@ -13,7 +13,7 @@
     "NumberOfInstances"
 ]
 
-chain = create_query_structuring_chain(document_content_description, content_attr, model = "llama3")
+chain = create_query_structuring_chain(document_content_description, content_attr, model = "llama3.1")
 
 from tenacity import retry, retry_if_exception_type, stop_after_attempt
 from langchain_community.query_constructors.chroma import ChromaTranslator

diff --git a/structured_query/llm_service_structured_query_utils.py b/structured_query/llm_service_structured_query_utils.py
@@ -1,6 +1,7 @@
 import json
 import sys
 from langchain_community.chat_models import ChatOllama
+# from langchain_ollama.llms import OllamaLLM
 
 from langchain.chains.query_constructor.base import (
     get_query_constructor_prompt,
@@ -16,7 +17,7 @@
 
 attribute_info = attribute_info[1:]
 
-def create_query_structuring_chain(document_content_description, content_attr, model = "llama3"):
+def create_query_structuring_chain(document_content_description, content_attr, model = "llama3.1"):
 
     prompt = get_query_constructor_prompt(
         document_contents=document_content_description,