fixed customllm

models/bloom.py

@@ -7,7 +7,7 @@ from dotenv import load_dotenv
 from llama_index import (Document, GPTVectorStoreIndex, LLMPredictor,
                          PromptHelper, ServiceContext, StorageContext,
                          load_index_from_storage)
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 from utils.customLLM import CustomLLM
 
@@ -27,8 +27,21 @@ num_output = 525
 chunk_overlap_ratio = 0.2
 prompt_helper = PromptHelper(context_window, num_output, chunk_overlap_ratio)
 
+# create a pipeline
+pl = pipeline(
+    model=model,
+    tokenizer=tokenizer,
+    task="text-generation",
+    # device=0, # GPU device number
+    # max_length=512,
+    do_sample=True,
+    top_p=0.95,
+    top_k=50,
+    temperature=0.7
+)
+
 # define llm
-llm_predictor = LLMPredictor(llm=CustomLLM(model, tokenizer))
+llm_predictor = LLMPredictor(llm=CustomLLM(model_pipeline=pl))
 service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
 
 def prepare_data(file_path:str):

utils/customLLM.py

@@ -1,24 +1,14 @@
 from typing import Any, List, Mapping, Optional
 
 from langchain.llms.base import LLM
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import Pipeline
 
 
 class CustomLLM(LLM):
-    
+    pipeline = None
     # Create the pipeline for question answering
-    def __init__(self, model: AutoModelForCausalLM, tokenizer: AutoTokenizer):
-        self.pipeline = pipeline(
-        model=model,
-        tokenizer=tokenizer,
-        task="text-generation",
-        # device=0, # GPU device number
-        # max_length=512,
-        do_sample=True,
-        top_p=0.95,
-        top_k=50,
-        temperature=0.7
-    )
+    def __init__(self, model_pipeline: Pipeline):
+        self.pipeline = model_pipeline
 
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         prompt_length = len(prompt)