update llama custom

models/llamaCustom.py

@@ -38,10 +38,10 @@ CHUNK_OVERLAP_RATION = 0.2
 
 
 @st.cache_resource
-def load_model(mode_name: str):
+def load_model(model_name: str):
     # llm_model_name = "bigscience/bloom-560m"
-    tokenizer = AutoTokenizer.from_pretrained(mode_name)
-    model = AutoModelForCausalLM.from_pretrained(mode_name, config="T5Config")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, config="T5Config")
 
     pipe = pipeline(
         task="text-generation",
@@ -62,11 +62,11 @@ class CustomLLM(LLM):
     llm_model_name: str
     pipeline: Any
 
-    def __init__(self, model_name: str):
-        # super().__init__()
+    def __init__(self, llm_model_name: str):
+        super().__init__()
         
-        self.llm_model_name = model_name
-        self.pipeline = load_model(mode_name=model_name)
+        self.llm_model_name = llm_model_name
+        self.pipeline = load_model(mode_name=llm_model_name)
 
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         prompt_length = len(prompt)