Update main.py

main.py

@@ -435,19 +435,18 @@ async def start():
     from transformers import pipeline
     from langchain_huggingface.llms import HuggingFacePipeline
     repo_id = "meta-llama/Llama-3.2-3B-Instruct"
-    pipe = pipeline(
-        "text-generation",
-        model=repo_id,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-    )
     #model = HuggingFaceEndpoint(
     #    repo_id=repo_id,
     #    max_new_tokens=6000,
     #    temperature=1.0,
     #    streaming=True
     #)
-    model = HuggingFacePipeline(pipeline=pipe)
+    model = HuggingFacePipeline.from_model_id(
+        model_id=repo_id,
+        task="text-generation",
+        device_map="auto",
+        pipeline_kwargs={"max_new_tokens": 6000},
+    )
     if not cl.user_session.get("saveMemory"):
         cl.user_session.set("saveMemory", "")
     cl.user_session.set("memory", ConversationBufferMemory(return_messages=True))