Hugging Face's logo

Spaces:
numind
/
NuExtract-v1.5
Running on L4

App Files Community
NuExtract-v1.5
File size: 6,515 Bytes 
b99458b
 
 
8103be7
b99458b
 
 
8103be7
 
b99458b
8103be7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e9ce3cc
8103be7
 
 
e9ce3cc
 
8103be7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
 
import os
import re
import time
import json
from itertools import cycle

import torch
import gradio as gr
from urllib.parse import unquote 
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList

from data import extract_leaves, split_document, handle_broken_output, clean_json_text, sync_empty_fields
from examples import examples as input_examples
from nuextract_logging import log_event


MAX_INPUT_SIZE = 10_000
MAX_NEW_TOKENS = 4_000
MAX_WINDOW_SIZE = 4_000

markdown_description = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>NuExtract</title>
</head>
<body>
    <img src="https://cdn.prod.website-files.com/638364a4e52e440048a9529c/64188f405afcf42d0b85b926_logo_numind_final.png" alt="NuMind Logo" style="vertical-align: middle;width: 200px; height: 50px;"> 
    <p>We are a startup developing custom information extraction models. NuExtract is a zero-shot model.</p>
    <p>If you want the best performance on your problem, please contact us :).</p> 
    <br>
    <ul>
        <li><strong>Webpage</strong>: <a href="https://www.numind.ai/">https://www.numind.ai/</a></li>
    </ul>
    <br>
    <h1>NuExtract-v1.5</h1>
    <p>NuExtract-v1.5 is a fine-tuning of Phi-3.5-mini-instruct, trained on a private high-quality dataset for structured information extraction. 
    It supports long documents and several languages (English, French, Spanish, German, Portuguese, and Italian). 
    To use the model, provide an input text and a JSON template describing the information you need to extract.</p>
    <ul>
        <li><strong>Model</strong>: <a href="https://huggingface.co/numind/NuExtract-v1.5">numind/NuExtract-v1.5</a></li>
    </ul>
    <i>NOTE: in this space we restrict the model inputs to a maximum length of 10k tokens, with anything over 4k being processed in a sliding window. For full model performance, self-host the model or contact us.</i>
</body>
</html>
"""


def highlight_words(input_text, json_output):
    colors = cycle(["#90ee90", "#add8e6", "#ffb6c1", "#ffff99", "#ffa07a", "#20b2aa", "#87cefa", "#b0e0e6", "#dda0dd", "#ffdead"])
    color_map = {}
    highlighted_text = input_text

    leaves = extract_leaves(json_output)
    for path, value in leaves:
        path_key = tuple(path)
        if path_key not in color_map:
            color_map[path_key] = next(colors)
        color = color_map[path_key]
        # highlighted_text = highlighted_text.replace(f" {value}", f" <span style='background-color: {color};'>{unquote(f'{value}')}</span>")

        pattern = rf"( |\n|\t){value}( |\n|\t)"
        replacement = f" <span style='background-color: {color};'>{unquote(value)}</span> "
        highlighted_text = re.sub(pattern, replacement, highlighted_text, flags=re.IGNORECASE)

    return highlighted_text

def predict_chunk(text, template, current, model, tokenizer):
    current = clean_json_text(current)

    input_llm =  f"<|input|>\n### Template:\n{template}\n### Current:\n{current}\n### Text:\n{text}\n\n<|output|>" + "{"
    input_ids = tokenizer(input_llm, return_tensors="pt", truncation=True, max_length=MAX_INPUT_SIZE).to("cuda")
    output = tokenizer.decode(model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS)[0], skip_special_tokens=True)

    return clean_json_text(output.split("<|output|>")[1])

def sliding_window_prediction(template, text, model, tokenizer, window_size=4000, overlap=128):
    # Split text into chunks of n tokens
    tokens = tokenizer.tokenize(text)
    chunks = split_document(text, window_size, overlap, tokenizer)

    # Iterate over text chunks
    prev = template
    full_pred = ""
    
    for i, chunk in enumerate(chunks):
        print(f"Processing chunk {i}...")
        pred = predict_chunk(chunk, template, prev, model, tokenizer)

        # Handle broken output
        pred = handle_broken_output(pred, prev)
        
        # create highlighted text
        highlighted_pred = highlight_words(text, json.loads(pred))

        # Sync empty fields
        synced_pred = sync_empty_fields(json.loads(pred), json.loads(template))
        synced_pred = json.dumps(synced_pred, indent=4)

        # Return progress, current prediction, and updated HTML
        yield f"Processed chunk {i+1}/{len(chunks)}", synced_pred, highlighted_pred

        # Iterate
        prev = pred


######

# Load the model and tokenizer
model_name = "numind/NuExtract-v1.5"
auth_token = os.environ.get("HF_TOKEN") or True
model = AutoModelForCausalLM.from_pretrained(model_name, 
                                             trust_remote_code=True, 
                                             torch_dtype=torch.bfloat16,
                                             device_map="auto", use_auth_token=auth_token)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=auth_token)
model.eval()

def gradio_interface_function(template, text):
    # reject invalid JSON
    try:
        template_json = json.loads(template)
    except:
        yield "", "Invalid JSON template", ""
        return  # End the function since there was an error

    if len(tokenizer.tokenize(text)) > MAX_INPUT_SIZE:
        yield "", "Input text too long for space. Download model to use unrestricted.", ""
        return  # End the function since there was an error

    # Initialize the sliding window prediction process
    prediction_generator = sliding_window_prediction(template, text, model, tokenizer, window_size=MAX_WINDOW_SIZE)

    # Iterate over the generator to return values at each step
    for progress, full_pred, html_content in prediction_generator:
        # yield gr.update(value=chunk_info), gr.update(value=progress), gr.update(value=full_pred), gr.update(value=html_content)
        yield progress, full_pred, html_content

    log_event(text, template, full_pred)

# Set up the Gradio interface
iface = gr.Interface(
    description=markdown_description,
    fn=gradio_interface_function,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter Template here...", label="Template"),
        gr.Textbox(lines=2, placeholder="Enter input Text here...", label="Input Text"),
    ],
    outputs=[
        gr.Textbox(label="Progress"),
        gr.Textbox(label="Model Output"),
        gr.HTML(label="Model Output with Highlighted Words"),
    ],
    examples=input_examples,
    # live=True  # Enable real-time updates
)

iface.launch(debug=True, share=True)