Update main.py

main.py

@@ -431,14 +431,23 @@ async def start():
     
     os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HUGGINGFACEHUB_API_TOKEN']
     #repo_id = "mistralai/Mistral-7B-Instruct-v0.3"
+    import torch
+    from transformers import pipeline
+    from langchain_huggingface.llms import HuggingFacePipeline
     repo_id = "meta-llama/Llama-3.2-3B-Instruct"
-
-    model = HuggingFaceEndpoint(
-        repo_id=repo_id,
-        max_new_tokens=6000,
-        temperature=1.0,
-        streaming=True
+    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)
     if not cl.user_session.get("saveMemory"):
         cl.user_session.set("saveMemory", "")
     cl.user_session.set("memory", ConversationBufferMemory(return_messages=True))