chat_model over llm

app/chains.py

@@ -1,8 +1,7 @@
 import os
 from dotenv import load_dotenv
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from langchain_core.runnables import RunnablePassthrough
-from transformers import AutoTokenizer
 import schemas
 from prompts import (
     raw_prompt,
@@ -16,17 +15,17 @@ load_dotenv()
 
 MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
 llm = HuggingFaceEndpoint(
     model=MODEL_ID,
     huggingfacehub_api_token=os.environ['HF_TOKEN'],
     max_new_tokens=512,
-    stop_sequences=[tokenizer.eos_token],
+    stop_sequences=["[EOS]", "<|end_of_text|>"],
     streaming=True,
 )
 
-simple_chain = (raw_prompt | llm).with_types(input_type=schemas.UserQuestion)
+chat_model = ChatHuggingFace(llm=llm)
+
+simple_chain = (raw_prompt | chat_model).with_types(input_type=schemas.UserQuestion)
 
 # # TODO: create formatted_chain by piping raw_prompt_formatted and the LLM endpoint.
 # formatted_chain = None