Spaces:
Sleeping
Sleeping
from typing import Dict, Union | |
from gliner import GLiNER | |
import gradio as gr | |
model = GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5").to('cpu') | |
text1 = """ | |
"I recently purchased the Sony WH-1000XM4 Wireless Noise-Canceling Headphones from Amazon and I must say, I'm thoroughly impressed. The package arrived in New York within 2 days, thanks to Amazon Prime's expedited shipping. | |
The headphones themselves are remarkable. The noise-canceling feature works like a charm in the bustling city environment, and the 30-hour battery life means I don't have to charge them every day. Connecting them to my Samsung Galaxy S21 was a breeze, and the sound quality is second to none. | |
I also appreciated the customer service from Amazon when I had a question about the warranty. They responded within an hour and provided all the information I needed. | |
However, the headphones did not come with a hard case, which was listed in the product description. I contacted Amazon, and they offered a 10% discount on my next purchase as an apology. | |
Overall, I'd give these headphones a 4.5/5 rating and highly recommend them to anyone looking for top-notch quality in both product and service.""" | |
open_ie_examples = [ | |
[ | |
f"Extract all brands, please", | |
text1, | |
0.5, | |
False | |
]] | |
def merge_entities(entities): | |
if not entities: | |
return [] | |
merged = [] | |
current = entities[0] | |
for next_entity in entities[1:]: | |
if next_entity['entity'] == current['entity'] and (next_entity['start'] == current['end'] + 1 or next_entity['start'] == current['end']): | |
current['word'] += ' ' + next_entity['word'] | |
current['end'] = next_entity['end'] | |
else: | |
merged.append(current) | |
current = next_entity | |
merged.append(current) | |
return merged | |
def process( | |
prompt:str, text, threshold: float, nested_ner: bool, labels: str = ["match"] | |
) -> Dict[str, Union[str, int, float]]: | |
text = prompt + "\n" + text | |
r = { | |
"text": text, | |
"entities": [ | |
{ | |
"entity": entity["label"], | |
"word": entity["text"], | |
"start": entity["start"], | |
"end": entity["end"], | |
"score": 0, | |
} | |
for entity in model.predict_entities( | |
text, labels, flat_ner=not nested_ner, threshold=threshold | |
) | |
], | |
} | |
r["entities"] = merge_entities(r["entities"]) | |
return r | |
with gr.Blocks(title="Open Information Extracting") as open_ie_interface: | |
prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here") | |
input_text = gr.Textbox(label="Text input", placeholder="Enter your text here") | |
threshold = gr.Slider(0, 1, value=0.3, step=0.01, label="Threshold", info="Lower the threshold to increase how many entities get predicted.") | |
nested_ner = gr.Checkbox(label="Nested NER", info="Allow for nested NER?") | |
output = gr.HighlightedText(label="Predicted Entities") | |
submit_btn = gr.Button("Submit") | |
theme=gr.themes.Base() | |
input_text.submit(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output) | |
prompt.submit(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output) | |
threshold.release(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output) | |
submit_btn.click(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output) | |
nested_ner.change(fn=process, inputs=[prompt, input_text, threshold, nested_ner], outputs=output) | |
if __name__ == "__main__": | |
open_ie_interface.launch() |