update llamaCustom

models/llamaCustom.py

@@ -20,10 +20,11 @@ from llama_index import (
     StorageContext,
     load_index_from_storage,
 )
+from pydantic import BaseModel
 from llama_index.llms.base import llm_completion_callback
 from llama_index.llms import CompletionResponse, CustomLLM, LLMMetadata
 from llama_index.prompts import Prompt
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, Pipeline
 
 from assets.prompts import custom_prompts
 
@@ -65,11 +66,13 @@ def load_model(model_name: str):
     return pipe
 
 
-class OurLLM(CustomLLM):
-    def __init__(self, model_name: str, pipeline):
-        super().__init__()  # Call the __init__ method of CustomLLM
-        self.model_name = model_name
-        self.pipeline = pipeline
+class OurLLM(CustomLLM, BaseModel):
+    # def __init__(self, model_name: str, pipeline):
+    #     super().__init__()  # Call the __init__ method of CustomLLM
+    #     self.model_name = model_name
+    #     self.pipeline = pipeline
+    model_name: str = ""
+    pipeline: Pipeline = None
 
     @property
     def metadata(self) -> LLMMetadata: