Update generator.py

generator.py

@@ -1,26 +1,30 @@
 import os
 from langchain_community.llms import HuggingFaceEndpoint
-
-# Load the API token from environment variable
-api_token = os.getenv("HUGGINGFACE_API_TOKEN")
+from pydantic import ValidationError
 
 def load_llm(repo_id="mistralai/Mistral-7B-Instruct-v0.2"):
     '''
     Load the LLM from the HuggingFace model hub
-
     Args:
         repo_id (str): The HuggingFace model ID
-
     Returns:
         llm (HuggingFaceEndpoint): The LLM model
     '''
 
-    repo_id = repo_id
-
-    llm = HuggingFaceEndpoint(
-        repo_id=repo_id, max_length=128, temperature=0.2, api_token=api_token)
-
-    return llm
+    try:
+        api_token = os.getenv("HUGGINGFACE_API_TOKEN")
+        llm = HuggingFaceEndpoint(
+            repo_id=repo_id, max_length=128, temperature=0.2, api_token=api_token)
+        return llm
+    except ValidationError as e:
+        print("Validation Error:", e)
+        # Log or handle the validation error appropriately
+        return None
+    except Exception as e:
+        print("Error:", e)
+        # Log or handle other exceptions
+        return None
 
 def guardrails():
     return None
+