File size: 2,663 Bytes
7ddeba9 a813767 7ddeba9 671cbf4 7ddeba9 671cbf4 7ddeba9 671cbf4 7ddeba9 671cbf4 7ddeba9 671cbf4 7ddeba9 671cbf4 |
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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
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
from PIL import Image
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):
# Process the uploaded PDF file
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 using OCR
image = Image.open(uploaded_image)
text = pytesseract.image_to_string(image)
st.write("Extracted text from the image:")
st.write(text)
# Apply entity extraction logic to the extracted text
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 {text}
Output: entity: type
"""
prompt_template = PromptTemplate(input_variables=["text"], template=template)
chain = LLMChain(llm=llm, prompt=prompt_template)
result = chain.run(text)
st.write("Extracted entities:")
entities = result.strip().split("\n")
table_data = [line.split(":") for line in entities]
st.table(table_data)
if __name__ == "__main__":
main() |