baLLseM / model /tools /predictor.py
hqms's picture
initial commit
8578816
raw
history blame contribute delete
977 Bytes
from . import Tool, LLMChain, llm
from langchain.memory import ConversationBufferMemory
from langchain_core.messages import SystemMessage
from langchain_core.prompts.chat import (
ChatPromptTemplate,
HumanMessagePromptTemplate,
MessagesPlaceholder,
)
system_message = "You are AI Assistant. You need to give crystal clear answer.\n"
template_messages = [
SystemMessage(content=system_message),
MessagesPlaceholder(variable_name="chat_history"),
HumanMessagePromptTemplate.from_template("{text}"),
]
prompt_template = ChatPromptTemplate.from_messages(template_messages)
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
word_problem_chain = LLMChain(llm=llm, prompt=prompt_template, memory=memory)
word_problem_tool = Tool.from_function(name="Reasoning Tool", \
func=word_problem_chain.run, \
description="Useful for when you need to answer logic-based/reasoning \
questions.",
)