app.py

import torch
from PIL import Image
import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import os
from threading import Thread

import pymupdf
import docx
from pptx import Presentation


MODEL_LIST = ["THUDM/glm-4v-9b"]
HF_TOKEN = os.environ.get("HF_TOKEN", None)
MODEL_ID = os.environ.get("MODEL_ID")
MODEL_NAME = MODEL_ID.split("/")[-1]

TITLE = "<h1>Multimodal Model for Complex Doc Extraction</h1>"

DESCRIPTION = f"""
<center>
<p>😊 A Demo For  Complex Doc Extraction via GLM4.
<br>
🚀 MODEL NOW: <a href="https://hf.co/{MODEL_ID}">{MODEL_NAME}</a>
<br>
✨ Important: Do not upload any sensitive documents.
<br>
🙇‍♂️ May be rebuilding from time to time.</p>
</center>"""

CSS = """
h1 {
    text-align: center;
    display: block;
}
"""

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(0)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model.eval()


def extract_text(path):
    return open(path, 'r').read()

def extract_pdf(path):
    doc = pymupdf.open(path)
    text = ""
    for page in doc:
        text += page.get_text()
    return text

def extract_docx(path):
    doc = docx.Document(path)
    data = []
    for paragraph in doc.paragraphs:
        data.append(paragraph.text)
    content = '\n\n'.join(data)
    return content

def extract_pptx(path):
    prs = Presentation(path)
    text = ""
    for slide in prs.slides:
        for shape in slide.shapes:
            if hasattr(shape, "text"):
                text += shape.text + "\n"
    return text

def mode_load(path):
    choice = ""
    file_type = path.split(".")[-1]
    print(file_type)
    if file_type in ["pdf", "txt", "py", "docx", "pptx", "json", "cpp", "md"]:
        if file_type.endswith("pdf"):
            content = extract_pdf(path)
        elif file_type.endswith("docx"):
            content = extract_docx(path)
        elif file_type.endswith("pptx"):
            content = extract_pptx(path)
        else:
            content = extract_text(path)
        choice = "doc"
        print(content[:100])
        return choice, content[:5000]
    elif file_type in ["png", "jpg", "jpeg", "bmp", "tiff", "webp"]:
        content = Image.open(path).convert('RGB')
        choice = "image"
        return choice, content
    else:
        raise gr.Error("Oops, unsupported files.")

@spaces.GPU()
def stream_chat(message, history: list, temperature: float, max_length: int, top_p: float, top_k: int, penalty: float):
    print(f'message is - {message}')
    print(f'history is - {history}')
    conversation = []
    prompt_files = []
    if message["files"]:
        choice, contents = mode_load(message["files"][-1])
        if choice == "image":
            conversation.append({"role": "user", "image": contents, "content": message['text']})
        elif choice == "doc":
            format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
            conversation.append({"role": "user", "content": format_msg})
    else:
        if len(history) == 0:
            #raise gr.Error("Please upload an image first.")
            contents = None
            conversation.append({"role": "user", "content": message['text']})
        else:
            #image = Image.open(history[0][0][0])
            for prompt, answer in history:
                if answer is None:
                    prompt_files.append(prompt[0])
                    conversation.extend([{"role": "user", "content": ""},{"role": "assistant", "content": ""}])
                else:
                    conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
            choice, contents = mode_load(prompt_files[-1])
            if choice == "image":
                conversation.append({"role": "user", "image": contents, "content": message['text']})
            elif choice == "doc":
                format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
                conversation.append({"role": "user", "content": format_msg})
    print(f"Conversation is -\n{conversation}")

    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)

    generate_kwargs = dict(
        max_length=max_length,
        streamer=streamer,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        repetition_penalty=penalty,
        eos_token_id=[151329, 151336, 151338],
    )
    gen_kwargs = {**input_ids, **generate_kwargs}

    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=gen_kwargs)
        thread.start()
        buffer = ""
        for new_text in streamer:
            buffer += new_text
            yield buffer
 



chatbot = gr.Chatbot()
chat_input = gr.MultimodalTextbox(
    interactive=True,
    placeholder="Enter message or upload a file one time...",
    show_label=False,

) 

prompt_extraction = '''
Please extract the Key information from the given document, including:
1 Date (DD/MM/YY)
2 Name(s) of Account Holder(s)
3 ID Document No.
4 ID Document Type
5 (Checked?)Are you the existing Personal e-Banking user?(expected result: yes/no)
6 (Checked?)Delete all Registered Bill Account Numbers for bill payment via any accounts under the above-mentioned ID Document or Business Registration No...(expected result: yes/no)
7 (Checked?)Add the following Bill Account Numbers for bill payment via any accounts under the above-mentioned ID Document or Business Registration No...(expected result: yes/no)
8 Signature(s) of Account Holder(s)
9 Account No.

Please Notice:
Return in JSON format, plain(no nested structure)
9 keys as listed in Objective section, don't ignore any item from the list even the value is None.
Please note:
Your accuracy of detecting checkboxes is relatively low, and unchecked boxes are often mistakenly identified as checked/true (resulting in many false positives).
Therefore, for all fields marked "(Checked?)", please be extra cautious and carefully examine the image before extracting the information, and if you are unsure, please default to false.
'''
prompt_signature = '''
I need your help to compare and score the similarity of two signatures. You should evaluate the signatures based on several dimensions and calculate a final similarity score. Here are the detailed instructions and dimensions for the comparison:

1. **General Shape and Flow**:
   - Compare the overall shape and flow of the two signatures.
   - Score from 0 to 10, where 0 means completely different and 10 means identical.

2. **Consistency of Loops and Strokes**:
   - Evaluate the presence and consistency of loops and strokes in the signatures.
   - Score from 0 to 10 based on the similarity of these features.

3. **Signature Characteristics**:
   - Compare specific characteristics such as dots, dashes, and unique flourishes.
   - Score from 0 to 10 based on the presence and similarity of these unique features.

4. **Stroke Pressure and Line Thickness**:
   - Analyze the pressure and thickness of the lines in the signatures.
   - Score from 0 to 10 based on how similar the pressure and thickness are between the two signatures.

5. **Angle and Slope**:
   - Evaluate the angle and slope of the characters in the signatures.
   - Score from 0 to 10 based on how similar the angles and slopes are.

6. **Spacing and Proportions**:
   - Compare the spacing between characters and the proportions of the signatures.
   - Score from 0 to 10 based on the similarity of spacing and proportions.

After scoring each dimension, calculate the final similarity score by averaging the scores from all dimensions. The final similarity score should be a value between 0 and 10, where 0 indicates no similarity and 10 indicates identical signatures.

Here is an example output format for the comparison:

```
General Shape and Flow: 8
Consistency of Loops and Strokes: 7
Signature Characteristics: 6
Stroke Pressure and Line Thickness: 5
Angle and Slope: 8
Spacing and Proportions: 7

Final Similarity Score: 6.83
```

Please help with this comparison and scoring for the two provided signatures. 

Important: Only score based on visual intuition, do not run code or provide coding solutions
'''




EXAMPLES = [
        [{"text": prompt_extraction, "files": ["./IMG_2700.png"]}],
        [{"text": prompt_signature , "files": ["./2.jpg"]}],
        [{"text": "Is it real?", "files": ["./spacecat.png"]}]
]

with gr.Blocks(css=CSS, theme="soft",fill_height=True) as demo:
    gr.HTML(TITLE)
    gr.HTML(DESCRIPTION)
    gr.ChatInterface(
        fn=stream_chat,
        multimodal=True,
        textbox=chat_input,
        chatbot=chatbot,
        fill_height=True,
        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Slider(
                minimum=0,
                maximum=1,
                step=0.1,
                value=0.8,
                label="Temperature",
                render=False,
            ),
            gr.Slider(
                minimum=1024,
                maximum=8192,
                step=1,
                value=4096,
                label="Max Length",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.1,
                value=1.0,
                label="top_p",
                render=False,
            ),
            gr.Slider(
                minimum=1,
                maximum=20,
                step=1,
                value=10,
                label="top_k",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=2.0,
                step=0.1,
                value=1.0,
                label="Repetition penalty",
                render=False,
            ),
        ],
    ),
    gr.Examples(EXAMPLES,[chat_input])


if __name__ == "__main__":
    demo.queue(api_open=False).launch(show_api=False, share=False)