updated code

models/llamaCustom.py

@@ -20,6 +20,7 @@ from llama_index import (
     StorageContext,
     load_index_from_storage,
 )
+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
@@ -78,6 +79,7 @@ class OurLLM(CustomLLM):
             model_name=self.model_name,
         )
 
+    @llm_completion_callback()
     def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         prompt_length = len(prompt)
         response = self.pipeline(prompt, max_new_tokens=NUM_OUTPUT)[0]["generated_text"]
@@ -86,6 +88,7 @@ class OurLLM(CustomLLM):
         text = response[prompt_length:]
         return CompletionResponse(text=text)
 
+    @llm_completion_callback()
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         raise NotImplementedError()
 

requirements.txt

@@ -1,4 +1,4 @@
-llama_index==0.7.16
+llama_index
 torch
 transformers
 panda