fix: add types

load_models.py

@@ -14,7 +14,6 @@ from transformers import (
     LlamaTokenizer,
     GenerationConfig,
     pipeline,
-    TextStreamer
 )
 
 

main.py

@@ -31,7 +31,6 @@ class Predict(BaseModel):
 class Delete(BaseModel):
     filename: str
 
-
 class MyCustomAsyncHandler(AsyncCallbackHandler):
     def on_llm_new_token(self, token: str, **kwargs) -> None:
         print(f" token: {token}")
@@ -66,7 +65,7 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks = [MyCustomAsyncHandler])
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks=[MyCustomAsyncHandler()])
 
 template = """you are a helpful, respectful and honest assistant. When answering questions, you should only use the documents provided.
 You should only answer the topics that appear in these documents.