Update main.py

main.py

@@ -431,22 +431,18 @@ 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.1-8B-Instruct"
     #model = HuggingFaceEndpoint(
     #    repo_id=repo_id,
     #    max_new_tokens=6000,
     #    temperature=1.0,
     #    streaming=True
     #)
-    model = HuggingFacePipeline.from_model_id(
-        model_id=repo_id,
-        task="text-generation",
-        device_map="auto",
-        pipeline_kwargs={"max_new_tokens": 6000},
-    )
+    from transformers import LlamaTokenizerFast, MistralForCausalLM
+    import torch
+    tokenizer = LlamaTokenizerFast.from_pretrained('mistralai/Mistral-Small-Instruct-2409')
+    tokenizer.pad_token = tokenizer.eos_token
+
+    model = MistralForCausalLM.from_pretrained('mistralai/Mistral-Small-Instruct-2409', torch_dtype=torch.bfloat16)
     if not cl.user_session.get("saveMemory"):
         cl.user_session.set("saveMemory", "")
     cl.user_session.set("memory", ConversationBufferMemory(return_messages=True))