feat: update rag example to use parameters (#122)

empirical-run · Apr 12, 2024 · 91112b7 · 91112b7
1 parent 3970f9c
commit 91112b7
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/examples/rag/empiricalrc.json b/examples/rag/empiricalrc.json
@@ -2,9 +2,24 @@
   "$schema": "https://assets.empirical.run/config/schema/v0.4.0.json",
   "runs": [
     {
-      "name": "rag script run",
       "type": "py-script",
       "path": "rag.py",
+      "parameters": {
+        "model": "gpt-3.5-turbo"
+      },
+      "scorers": [
+        {
+          "type": "py-script",
+          "path": "score.py"
+        }
+      ]
+    },
+    {
+      "type": "py-script",
+      "path": "rag.py",
+      "parameters": {
+        "model": "gpt-4-turbo-preview"
+      },
       "scorers": [
         {
           "type": "py-script",

diff --git a/examples/rag/rag.py b/examples/rag/rag.py
@@ -1,10 +1,7 @@
-import nest_asyncio
 from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
 from llama_index.llms.openai import OpenAI as LlamaOpenAI
 from llama_index.core.node_parser import SentenceSplitter
 
-nest_asyncio.apply()
-
 
 def execute(inputs, parameters):
     # load documents
@@ -21,7 +18,9 @@ def build_query_engine(llm):
         query_engine = vector_index.as_query_engine(similarity_top_k=2, llm=llm)
         return query_engine
 
-    query_engine = build_query_engine(llm=LlamaOpenAI(model="gpt-3.5-turbo"))
+    query_engine = build_query_engine(
+        llm=LlamaOpenAI(model=parameters.get("model", "gpt-3.5-turbo"))
+    )
     response = query_engine.query(question)
     output = response.response
     contexts = [c.node.get_content() for c in response.source_nodes]