File size: 1,869 Bytes
7ddeba9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import streamlit as st
import pytesseract
from tempfile import NamedTemporaryFile
from langchain.document_loaders import PyPDFLoader
from langchain.llms import CTransformers
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
def main():
st.title("Invoice Entity Extractor 📚")
uploaded_file = st.sidebar.file_uploader("Upload a PDF file", type="pdf")
uploaded_image = st.sidebar.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
if uploaded_file is not None:
process_pdf(uploaded_file)
elif uploaded_image is not None:
process_image(uploaded_image)
def process_pdf(uploaded_file):
# Save the uploaded file to a temporary location
with NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(uploaded_file.read())
temp_file_path = temp_file.name
loader = PyPDFLoader(temp_file_path)
pages = loader.load()
st.write(f"Number of pages: {len(pages)}")
for page in pages:
st.write(page.page_content)
llm = CTransformers(model="llama-2-7b-chat.ggmlv3.q4_0.bin",model_type="llama",
config={'max_new_tokens':128,'temperature':0.01})
template = """Extract invoice number, name of organization, address, date,
Qty, Rate ,Tax ,Amount {pages}
Output : entity : type
"""
prompt_template = PromptTemplate(input_variables=["pages"], template=template)
chain = LLMChain(llm=llm, prompt=prompt_template)
result = chain.run(pages=pages[0].page_content)
st.write("Extracted entities:")
entities = result.strip().split("\n")
table_data = [line.split(":") for line in entities]
st.table(table_data)
def process_image(uploaded_image):
# Process the uploaded image
st.write("Image processing is not implemented yet.")
if __name__ == "__main__":
main()
|