ankur-bohra's picture
Add basic structure
0d99179
raw
history blame
1.68 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.6,
max_tokens=300,
n=1,
request_timeout=None,
model_kwargs={
'stop': None,
'top_p': 1,
}
)
# Build category chain
system_message_prompt = SystemMessagePromptTemplate.from_template(
"You are tasked with developing an OCR data extraction system for hotel bills in PDF "
"format given as text. The system should extract important information necessary for "
"the reimbursement process from a college. Your prompt should fetch the following "
"essential details from the hotel bill: hotel name, address, bill number/invoice "
"number, booking ID / confirmation ID / booking number, check-in date and time, "
"check-out date and time, total amount, booking platform, bill date.\n"
"Ensure that the system accurately extracts the above information from the OCR text "
"of the hotel bill.\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)