Upload 7 files

main.py

@@ -11,6 +11,8 @@ This script performs the following tasks:
 
 # --- Importing required libraries ---
 from zigistry import constants
+from llama_index.llms.openrouter import OpenRouter
+from llama_index.core.llms import ChatMessage
 from zigistry import pre_requisite
 import gradio as gr
 from llama_index.core import (
@@ -36,13 +38,12 @@ if not pre_requisite.CorrectlyCompletedPreRequisites():
 
 
 # --- Configure LLM and embedding model ---
-Settings.llm = HuggingFaceInferenceAPI(
+Settings.llm = OpenRouter(
     model_name=constants.LLM,
-    tokenizer_name=constants.LLM,
+    api_key=constants.OR_TOKEN,
     context_window=3000,
-    token=constants.HF_TOKEN,
-    max_new_tokens=512,
-    generate_kwargs={"temperature": constants.TEMPERATURE},
+    max_tokens=512,
+    temperature=constants.TEMPERATURE,
 )
 
 Settings.embed_model = HuggingFaceEmbedding(model_name=constants.EMBEDDER)

zigistry/constants.py

@@ -5,9 +5,9 @@ print("Loaded environment variables.")
 load_dotenv()
 
 PERSIST_DIR = "PERSIST_DIR"
-LLM = "google/gemma-1.1-7b-it"
+LLM = "meta-llama/llama-guard-3-8b"
 EMBEDDER = "BAAI/bge-small-en-v1.5"
-HF_TOKEN = os.getenv("HF_TOKEN")
+OR_TOKEN = os.getenv("OR_TOKEN")
 TEMPERATURE = 0.1
 FILES = ["Zigistry.pdf"]
 

zigistry/pre_requisite.py

@@ -11,7 +11,7 @@ def CorrectlyCompletedPreRequisites() -> bool:
     if not os.path.exists(constants.PERSIST_DIR):
         error(f"Directory {constants.PERSIST_DIR} does not exist")
         return False
-    if constants.HF_TOKEN == None:
+    if constants.OR_TOKEN == None:
         error("Hugging Face token not set")
         return False
     return True