Merge pull request #21 from openworm/development

Update tests

Author: pgleeson

openworm_ai/graphrag/GraphRAG_test.py

@@ -15,6 +15,7 @@
 from llama_index.core import VectorStoreIndex, get_response_synthesizer
 from llama_index.core.retrievers import VectorIndexRetriever
 from llama_index.core.query_engine import RetrieverQueryEngine
+from llama_index.core import Settings
 
 
 # one extra dep
@@ -27,16 +28,18 @@
 STORE_DIR = "store"
 SOURCE_DOCUMENT = "source document"
 
+Settings.chunk_size = 3000
+Settings.chunk_overlap = 50
+
 
 def create_store(model):
     OLLAMA_MODEL = model.replace("Ollama:", "") if model is not LLM_GPT4o else None
 
-    json_inputs = glob.glob("processed/json/*/*.json")
-    # print_(json_inputs)
+    json_inputs = glob.glob("processed/json/papers/*.json")
 
     documents = []
     for json_file in json_inputs:
-        print_("Adding %s" % json_file)
+        print_("Adding file to document store: %s" % json_file)
 
         with open(json_file, encoding="utf-8") as f:
             doc_model = json.load(f)
@@ -60,36 +63,33 @@ def create_store(model):
                 if len(all_text) == 0:
                     all_text = " "
                 # print_(f'---------------------\n{all_text}\n---------------------')
-                src_info = (
-                    f"WormAtlas Handbook: [{title}, Section {section}]({src_page})"
-                )
+                src_type = "Publication"
+                if "wormatlas" in json_file:
+                    src_type = "WormAtlas Handbook"
+                src_info = f"{src_type}: [{title}, Section {section}]({src_page})"
                 doc = Document(text=all_text, metadata={SOURCE_DOCUMENT: src_info})
                 documents.append(doc)
 
-    if "-test" in sys.argv:
-        print_("Finishing before section requiring OPENAI_API_KEY...")
-
-    else:
-        print_("Creating a vector store index for %s" % model)
+    print_("Creating a vector store index for %s" % model)
 
-        STORE_SUBFOLDER = ""
+    STORE_SUBFOLDER = ""
 
-        if OLLAMA_MODEL is not None:
-            ollama_embedding = OllamaEmbedding(
-                model_name=OLLAMA_MODEL,
-            )
-            STORE_SUBFOLDER = "/%s" % OLLAMA_MODEL.replace(":", "_")
+    if OLLAMA_MODEL is not None:
+        ollama_embedding = OllamaEmbedding(
+            model_name=OLLAMA_MODEL,
+        )
+        STORE_SUBFOLDER = "/%s" % OLLAMA_MODEL.replace(":", "_")
 
-            # create an index from the parsed markdown
-            index = VectorStoreIndex.from_documents(
-                documents, embed_model=ollama_embedding, show_progress=True
-            )
-        else:
-            index = VectorStoreIndex.from_documents(documents)
+        # create an index from the parsed markdown
+        index = VectorStoreIndex.from_documents(
+            documents, embed_model=ollama_embedding, show_progress=True
+        )
+    else:
+        index = VectorStoreIndex.from_documents(documents)
 
-        print_("Persisting vector store index")
+    print_("Persisting vector store index")
 
-        index.storage_context.persist(persist_dir=STORE_DIR + STORE_SUBFOLDER)
+    index.storage_context.persist(persist_dir=STORE_DIR + STORE_SUBFOLDER)
 
 
 def load_index(model):
@@ -147,7 +147,7 @@ def get_query_engine(index_reloaded, model, similarity_top_k=4):
 
     # create a query engine for the index
     if OLLAMA_MODEL is not None:
-        llm = Ollama(model=OLLAMA_MODEL)
+        llm = Ollama(model=OLLAMA_MODEL, request_timeout=60.0)
 
         ollama_embedding = OllamaEmbedding(
             model_name=OLLAMA_MODEL,
@@ -159,6 +159,7 @@ def get_query_engine(index_reloaded, model, similarity_top_k=4):
             refine_template=refine_template,
             embed_model=ollama_embedding,
         )
+        # print(dir(query_engine.retriever))
 
         query_engine.retriever.similarity_top_k = similarity_top_k
 
@@ -184,27 +185,10 @@ def get_query_engine(index_reloaded, model, similarity_top_k=4):
     return query_engine
 
 
-def process_query(response, model):
+def process_query(query, model, verbose=False):
+    print_("Processing query: %s" % query)
     response = query_engine.query(query)
 
-    """
-    import pprint as pp
-
-    print(type(response))
-    print(dir(response))
-
-    print("------")
-    pp.pprint(response.metadata)
-    print("------")
-
-    for sn in response.source_nodes:
-        print("  -- ")
-        print(f' - {sn.score}: {sn.metadata['source document']}')
-        pp.pprint(sn)
-    print("------")
-    pp.pprint(response.response)
-    print("------")"""
-
     response_text = str(response)
 
     if "<think>" in response_text:  # Give deepseek a fighting chance...
@@ -217,7 +201,14 @@ def process_query(response, model):
     cutoff = 0.2
     files_used = []
     for sn in response.source_nodes:
-        # print(sn)
+        if verbose:
+            print_("===================================")
+            # print(dir(sn))
+            print_(sn.metadata["source document"])
+            print_("-------")
+            print_("Length of selection below: %i" % len(sn.text))
+            print_(sn.text)
+
         sd = sn.metadata["source document"]
 
         if sd not in files_used:
@@ -244,10 +235,10 @@ def process_query(response, model):
 
     llm_ver = get_llm_from_argv(sys.argv)
 
-    if "-q" not in sys.argv:
-        create_store(llm_ver)
-
     if "-test" not in sys.argv:
+        if "-q" not in sys.argv:
+            create_store(llm_ver)
+
         index_reloaded = load_index(llm_ver)
         query_engine = get_query_engine(index_reloaded, llm_ver)
 
@@ -265,11 +256,18 @@ def process_query(response, model):
             "When was the first metazoan genome sequenced? Answer only with the year.","""
 
         queries = [
-            "The NeuroPAL transgene is amazing. Give me some examples of fluorophores in it.",
             "What is the main function of cell pair AVB?",
-            "What can you tell me about Alan Coulson?",
             "In what year was William Shakespeare born? ",
+            "Tell me about the egg laying apparatus in C. elegans",
+            "Tell me briefly about the neuronal control of C. elegans locomotion and the influence of monoamines.",
+            "What can you tell me about Alan Coulson?",
+            "The NeuroPAL transgene is amazing. Give me some examples of fluorophores in it.",
         ]
+        """queries = [
+            "What can you tell me about Alan Coulson?",
+        ]"""
+
+        print_("Processing %i queries" % len(queries))
 
         for query in queries:
             process_query(query, llm_ver)

openworm_ai/quiz/QuizMaster.py

@@ -47,7 +47,7 @@ def save_quiz(num_questions, num_answers, llm_ver, quiz_scope, temperature=0):
     last_question = None
 
     indexing = ["1", "2", "3", "4"]
-    
+
     for line in response.split("\n"):
         if len(line.strip()) > 0:
             if "QUESTION" in line or line[-1] == "?":
@@ -87,6 +87,7 @@ def save_quiz(num_questions, num_answers, llm_ver, quiz_scope, temperature=0):
         # quiz_json = "openworm_ai/quiz/samples/GPT4o_10questions.json"
 
         quiz_json = "openworm_ai/quiz/samples/GPT4o_100questions.json"
+        quiz_json = "openworm_ai/quiz/samples/GPT4o_100questions_celegans.json"
 
         quiz = MultipleChoiceQuiz.from_file(quiz_json)
 
@@ -122,9 +123,10 @@ def save_quiz(num_questions, num_answers, llm_ver, quiz_scope, temperature=0):
 
             from openworm_ai.utils.llms import ask_question_get_response
 
-            resp = ask_question_get_response(
+            orig_resp = ask_question_get_response(
                 full_question, llm_ver, print_question=False
             ).strip()
+            resp = orig_resp
 
             if "<think>" in resp:  # Give deepseek a fighting chance...
                 resp = (
@@ -133,6 +135,8 @@ def save_quiz(num_questions, num_answers, llm_ver, quiz_scope, temperature=0):
                 resp = resp.replace("\n", " ").strip()
                 guess = resp[-1]
             else:
+                if "\n" in resp:
+                    resp = resp.split("\n")[0]
                 guess = resp.split(":")[0].strip()
                 if " " in guess:
                     guess = guess[0]
@@ -143,7 +147,7 @@ def save_quiz(num_questions, num_answers, llm_ver, quiz_scope, temperature=0):
             if guess in presented_answers:
                 g = presented_answers[guess]
             else:
-                g = "%s (cannot be interpreted!)" % guess
+                g = "[%s] [[%s]] (this cannot be interpreted!)" % (guess, orig_resp)
             print(
                 f" >> {qi}) Is their guess of ({g}) for ({q}) correct (right answer: {correct_text})? {correct_guess}"
             )

openworm_ai/quiz/quiz_all.py

@@ -2,25 +2,29 @@
 import time
 import random
 import datetime
+
 from openworm_ai.utils.llms import (
     LLM_OLLAMA_LLAMA32_1B,
-    LLM_GPT4o,
-    LLM_GEMINI,
-    LLM_CLAUDE37,
-    LLM_GPT35,
-    LLM_OLLAMA_PHI4,
-    LLM_OLLAMA_GEMMA2,
-    LLM_OLLAMA_GEMMA,
-    LLM_OLLAMA_QWEN,
-    LLM_OLLAMA_TINYLLAMA,
+    LLM_OLLAMA_LLAMA32_3B,
+    # LLM_GPT4o,
+    # LLM_GEMINI,
+    # LLM_CLAUDE37,
+    # LLM_GPT35,
+    # LLM_OLLAMA_PHI4,
+    # LLM_OLLAMA_GEMMA2,
+    # LLM_OLLAMA_GEMMA,
+    # LLM_OLLAMA_QWEN,
+    # LLM_OLLAMA_TINYLLAMA,
     ask_question_get_response,
 )
+
+
 from openworm_ai.quiz.Templates import (
     ASK_Q,
 )  # Ensure this matches the correct import path
 
-iteration_per_day = 1
 field = "celegans"  # general/science/celegans
+iteration_per_day = 3
 current_date = datetime.datetime.now().strftime("%d-%m-%y")
 SOURCE_QUESTIONS_FILE = "openworm_ai/quiz/samples/GPT4o_100questions_celegans.json"
 OUTPUT_FILENAME = f"llm_scores_{field}_{current_date}_{iteration_per_day}.json"
@@ -35,16 +39,17 @@ def load_llms():
     """Loads only the selected LLMs: Ollama Llama3 and GPT-3.5."""
     llms = [
         LLM_OLLAMA_LLAMA32_1B,
-        LLM_GPT4o,
-        LLM_GEMINI,
-        LLM_CLAUDE37,
-        LLM_GPT35,
-        LLM_OLLAMA_PHI4,
-        LLM_OLLAMA_GEMMA2,
+        LLM_OLLAMA_LLAMA32_3B,
+        # LLM_GPT4o,
+        #####LLM_GEMINI,
+        ####LLM_CLAUDE37,
+        ###LLM_GPT35,
+        ##LLM_OLLAMA_PHI4,
+        # LLM_OLLAMA_GEMMA2,
         # LLM_OLLAMA_DEEPSEEK - unable to answer A-D(too few params?),
-        LLM_OLLAMA_GEMMA,
-        LLM_OLLAMA_QWEN,
-        LLM_OLLAMA_TINYLLAMA,
+        # LLM_OLLAMA_GEMMA,
+        # LLM_OLLAMA_QWEN,
+        # LLM_OLLAMA_TINYLLAMA,
         # LLM_OLLAMA_FALCON2 - 'only an assistant with no acess to external resources',
         # LLM_OLLAMA_CODELLAMA - understands only a fraction of questions, doesnt understand prompts
     ]  # Defined constants

openworm_ai/utils/llms.py

@@ -25,6 +25,10 @@
 LLM_CMD_LINE_ARGS["-o-l32"] = LLM_OLLAMA_LLAMA32
 LLM_OLLAMA_LLAMA32_1B = "Ollama:llama3.2:1b"
 LLM_CMD_LINE_ARGS["-o-l321b"] = LLM_OLLAMA_LLAMA32_1B
+
+LLM_OLLAMA_LLAMA32_3B = "Ollama:llama3.2:3b"
+LLM_CMD_LINE_ARGS["-o-l323b"] = LLM_OLLAMA_LLAMA32_3B
+
 LLM_OLLAMA_MISTRAL = "Ollama:mistral"
 LLM_CMD_LINE_ARGS["-o-m"] = LLM_OLLAMA_MISTRAL
 LLM_OLLAMA_TINYLLAMA = "Ollama:tinyllama"
@@ -45,6 +49,10 @@
 LLM_CMD_LINE_ARGS["-qw"] = LLM_OLLAMA_QWEN
 LLM_OLLAMA_CODELLAMA = "Ollama:codellama:latest"
 LLM_OLLAMA_FALCON2 = "Ollama:falcon2:latest"
+LLM_OLLAMA_FALCON2 = "Ollama:falcon2:latest"
+
+LLM_OLLAMA_OLMO2_7B = "Ollama:olmo2:7b"
+LLM_CMD_LINE_ARGS["-o-olmo27b"] = LLM_OLLAMA_OLMO2_7B
 
 OPENAI_LLMS = [LLM_GPT35, LLM_GPT4, LLM_GPT4o]
 
@@ -70,6 +78,7 @@
     LLM_OLLAMA_QWEN,
     LLM_OLLAMA_CODELLAMA,
     LLM_OLLAMA_FALCON2,
+    LLM_OLLAMA_OLMO2_7B,
 )
 
 
@@ -186,6 +195,7 @@ def get_llm(llm_ver, temperature):
 
     elif llm_ver in [
         LLM_OLLAMA_LLAMA32_1B,
+        LLM_OLLAMA_LLAMA32_3B,
         LLM_OLLAMA_MISTRAL,
         LLM_OLLAMA_TINYLLAMA,
         LLM_OLLAMA_PHI3,
@@ -197,6 +207,7 @@ def get_llm(llm_ver, temperature):
         LLM_OLLAMA_QWEN,
         LLM_OLLAMA_CODELLAMA,
         LLM_OLLAMA_FALCON2,
+        LLM_OLLAMA_OLMO2_7B,
     ]:
         from langchain_ollama.llms import OllamaLLM
 
@@ -336,6 +347,8 @@ def ask_question_get_response(
     import sys
 
     question = "What is the most common type of neuron in the brain?"
+    question = "Why is the worm C. elegans important to scientists?"
+    question = "Tell me briefly about the neuronal control of C. elegans locomotion and the influence of monoamines."
 
     llm_ver = get_llm_from_argv(sys.argv)