Merge pull request #24 from openworm/development

Improved testing

Author: pgleeson

corpus/papers/test/Wang2024_NeurotransmitterAtlas.pdf.json

@@ -1,5 +1,5 @@
 # Version of the Python module.
-__version__ = "0.2.6"
+__version__ = "0.2.7"
 
 
 def print_(msg, print_it=True):

openworm_ai/__init__.py

@@ -267,6 +267,7 @@ def process_query(query, model, verbose=False):
             "What are the main differences between NeuroML versions 1 and 2?",
             "What are the main types of cell in the C. elegans pharynx?",
             "Give me 3 facts about the coelomocyte system in C. elegans",
+            "Tell me about the neurotransmitter betaine in C. elegans",
         ]
 
         print_("Processing %i queries" % len(queries))

openworm_ai/graphrag/GraphRAG_test.py

@@ -83,6 +83,10 @@ def convert_to_json(paper_ref, paper_info, output_dir):
             "corpus/papers/test/SinhaEtAl2025.pdf.json",
             "https://elifesciences.org/articles/95135",
         ],
+        "Wang_et_al_2024": [
+            "corpus/papers/test/Wang2024_NeurotransmitterAtlas.pdf.json",
+            "https://elifesciences.org/articles/95402",
+        ],
     }
 
     # Loop through papers and process markdown sections

openworm_ai/parser/ParseLlamaIndexJson.py

@@ -1,24 +1,45 @@
 import json
 import os
+import sys
 import matplotlib.pyplot as plt
 import pandas as pd
 
+# ruff: noqa: F401
+from openworm_ai.utils.llms import (
+    LLM_OLLAMA_LLAMA32_1B,
+    LLM_OLLAMA_LLAMA32_3B,
+    LLM_GPT4o,
+    LLM_GEMINI_2F,
+    LLM_CLAUDE37,
+    LLM_GPT35,
+    LLM_OLLAMA_PHI4,
+    LLM_OLLAMA_GEMMA2,
+    LLM_OLLAMA_GEMMA,
+    LLM_OLLAMA_QWEN,
+    LLM_OLLAMA_TINYLLAMA,
+    ask_question_get_response,
+)
+
 # Define model parameters (LLM parameter sizes in billions)
 llm_parameters = {
+    LLM_GPT4o: 1760,
+    LLM_GPT35: 175,
     "GPT3.5": 20,
     "Phi4": 14,
     "Gemma2": 9,
     "Gemma": 7,
     "Qwen": 4,
     "Llama3.2": 1,
-    "TinyLlama":1.1,
+    "TinyLlama": 1.1,
     "GPT4o": 1760,
     "Gemini": 500,
-    "Claude 3.5 Sonnet": 175
+    "Claude 3.5 Sonnet": 175,
 }
 
 # Define model distributors for coloring
 model_distributors = {
+    LLM_GPT4o: "OpenAI",
+    LLM_GPT35: "OpenAI",
     "GPT3.5": "OpenAI",
     "GPT4o": "OpenAI",
     "Phi4": "Microsoft",
@@ -28,15 +49,15 @@
     "Claude 3.5 Sonnet": "Anthropic",
     "Qwen": "Alibaba",
     "Llama3.2": "Meta",
-    "TinyLlama":"Open Source"
+    "TinyLlama": "Open Source",
 }
 
 # Define quiz categories and corresponding file paths
 file_paths = {
-    #"General Knowledge": "openworm_ai/quiz/scores/general/llm_scores_general_24-02-25.json",
-    #"Science": "openworm_ai/quiz/scores/science/llm_scores_science_24-02-25.json",
-    #"C. Elegans": "openworm_ai/quiz/scores/celegans/llm_scores_celegans_24-02-25.json",
-    "RAG":"openworm_ai/quiz/scores/rag/llm_scores_rag_16-03-25_2.json"
+    # "General Knowledge": "openworm_ai/quiz/scores/general/llm_scores_general_24-02-25.json",
+    # "Science": "openworm_ai/quiz/scores/science/llm_scores_science_24-02-25.json",
+    # "C. Elegans": "openworm_ai/quiz/scores/celegans/llm_scores_celegans_24-02-25.json",
+    "RAG": "openworm_ai/quiz/scores/rag/llm_scores_rag_16-03-25_2.json"
 }
 
 # Folder to save figures
@@ -51,16 +72,19 @@
     "Microsoft": "purple",
     "Alibaba": "orange",
     "Meta": "cyan",
-    "Open Source":"yellow"
+    "Open Source": "yellow",
 }
 
 # Process each quiz category
 for category, file_path in file_paths.items():
-    save_path = os.path.join(figures_folder, f"llm_accuracy_vs_parameters_{category.replace(' ', '_').lower()}.png")
+    save_path = os.path.join(
+        figures_folder,
+        f"llm_accuracy_vs_parameters_{category.replace(' ', '_').lower()}.png",
+    )
 
     # Check if the file exists
     if not os.path.exists(file_path):
-        print(f"⚠️ Warning: File not found - {file_path}. Skipping this category.")
+        print(f"Warning: File not found - {file_path}. Skipping this category.")
         continue
 
     # Load JSON data
@@ -72,17 +96,19 @@
     for result in data.get("Results", []):  # Use .get() to avoid KeyError
         for key in llm_parameters:
             if key.lower() in result["LLM"].lower():
-                category_results.append({
-                    "Model": key,
-                    "Accuracy (%)": result["Accuracy (%)"],
-                    "Parameters (B)": llm_parameters[key],
-                    "Distributor": model_distributors.get(key, "Unknown")
-                })
+                category_results.append(
+                    {
+                        "Model": key,
+                        "Accuracy (%)": result["Accuracy (%)"],
+                        "Parameters (B)": llm_parameters[key],
+                        "Distributor": model_distributors.get(key, "Unknown"),
+                    }
+                )
                 break
 
     # Skip if no data
     if not category_results:
-        print(f"⚠️ No valid results found in {file_path}. Skipping...")
+        print(f"No valid results found in {file_path}. Skipping...")
         continue
 
     # Convert to DataFrame
@@ -94,11 +120,25 @@
     # Scatter plot with model labels, colored by distributor
     for distributor, color in distributor_colors.items():
         subset = df[df["Distributor"] == distributor]
-        plt.scatter(subset["Parameters (B)"], subset["Accuracy (%)"], s=100, color=color, label=distributor, edgecolor="black")
+        plt.scatter(
+            subset["Parameters (B)"],
+            subset["Accuracy (%)"],
+            s=100,
+            color=color,
+            label=distributor,
+            edgecolor="black",
+        )
 
     # Add model labels to each point
     for i, row in df.iterrows():
-        plt.text(row["Parameters (B)"], row["Accuracy (%)"], row["Model"], fontsize=10, ha="right", va="bottom")
+        plt.text(
+            row["Parameters (B)"],
+            row["Accuracy (%)"],
+            row["Model"],
+            fontsize=10,
+            ha="right",
+            va="bottom",
+        )
 
     # Log scale for x-axis (model parameters)
     plt.xscale("log")
@@ -113,5 +153,6 @@
     # Save figure
     plt.legend()
     plt.savefig(save_path)
-    print(f"✅ Saved plot: {save_path}")
-    plt.show()
+    print(f"Saved plot: {save_path}")
+    if "-nogui" not in sys.argv:
+        plt.show()