ankur-bohra's picture
Initial commit
317211f
raw
history blame
1.22 kB
from .model import InformationExtractedFromABillReceipt as PydanticModel
from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from langchain.prompts import (
ChatPromptTemplate,
HumanMessagePromptTemplate,
SystemMessagePromptTemplate,
)
model = ChatOpenAI(
temperature=0,
n=1,
model_kwargs= {
'stop': None,
'top_p': 1,
'frequency_penalty': 0,
'presence_penalty': 0,
}
)
# Build categorizing chain
system_message_prompt = SystemMessagePromptTemplate.from_template(
"You are an information extraction engine that outputs details from OCR processed "
"documents such as date/time/place of departure and arrival.\n"
"{format_instructions}"
)
human_message_prompt = HumanMessagePromptTemplate.from_template("{text}")
chat_prompt = ChatPromptTemplate.from_messages(
[system_message_prompt, human_message_prompt]
)
output_parser = PydanticOutputParser(pydantic_object=PydanticModel)
fixing_parser = OutputFixingParser.from_llm(llm=model, parser=output_parser)
chain = LLMChain(
llm=model, prompt=chat_prompt, output_parser=fixing_parser
)