Update main.py

main.py

@@ -437,9 +437,12 @@ async def start():
     #    temperature=1.0,
     #    streaming=True
     #)
-    from transformers import LlamaTokenizerFast, MistralForCausalLM
-    import torch
-    model = MistralForCausalLM.from_pretrained('mistralai/Mistral-Small-Instruct-2409', torch_dtype=torch.bfloat16)
+    from langchain_huggingface.llms import HuggingFacePipeline
+    model = HuggingFacePipeline.from_model_id(
+        model_id="mistralai/Mistral-Small-Instruct-2409",
+        task="text-generation",
+        pipeline_kwargs={"max_new_tokens": 10},
+    )
     if not cl.user_session.get("saveMemory"):
         cl.user_session.set("saveMemory", "")
     cl.user_session.set("memory", ConversationBufferMemory(return_messages=True))