Changes for AI21

requirements.txt

@@ -9,5 +9,6 @@ neuromllite
 matplotlib
 pylems
 pyneuroml
-paper-qa==4.0.0rc8
-langchain-google-genai
+paper-qa
+langchain-google-genai
+langchain_ai21

ring.py

@@ -16,7 +16,7 @@ LLM_AI21 = "AI21"
 
 OPENAI_LLMS = [LLM_GPT35, LLM_GPT4, LLM_GPT4o]
 
-PREF_ORDER_LLMS = (LLM_GEMINI, LLM_LLAMA2, LLM_GPT35, LLM_GPT4, LLM_GPT4o, AI21)
+PREF_ORDER_LLMS = (LLM_GEMINI, LLM_LLAMA2, LLM_GPT35, LLM_GPT4, LLM_GPT4o, LLM_AI21)
 
 
 def requires_openai_key(llm_ver):
@@ -101,6 +101,14 @@ def get_llm(llm_ver, temperature):
             model="gemini-pro", google_api_key=get_gemini_api_key()
         )
 
+    elif llm_ver == LLM_AI21:
+
+        from langchain_ai21 import AI21LLM
+
+        llm = AI21LLM(
+            model="j2-ultra"
+        )
+
     return llm
 
 

test.py

@@ -3,6 +3,7 @@ from ring import LLM_GPT4
 from ring import LLM_GPT4o
 from ring import LLM_LLAMA2
 from ring import LLM_GEMINI
+from ring import LLM_AI21
 from ring import PREF_ORDER_LLMS
 
 from ring import generate_response
@@ -16,6 +17,12 @@ llm_ver = LLM_GPT4o
 if "-g" in sys.argv:
     llm_ver = LLM_GEMINI
 
+if "-l" in sys.argv:
+    llm_ver = LLM_LLAMA2
+
+if "-a" in sys.argv:
+    llm_ver = LLM_AI21
+
 
 print("--------------------------------------------------------")
 print("Asking question:\n   %s" % question)