Roland Szabo commited on
Commit
c1d9469
0 Parent(s):

Initial commit

Browse files
Files changed (2) hide show
  1. main.py +41 -0
  2. requirements.txt +4 -0
main.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spacy
3
+ from transformers import pipeline
4
+ import boto3
5
+ nlp = spacy.load("en_core_web_sm")
6
+ ner_pipeline = pipeline("ner", model="Jean-Baptiste/roberta-large-ner-english", aggregation_strategy="simple", grouped_entities=True)
7
+
8
+
9
+ def greet(model_type, text):
10
+ if model_type == "Spacy":
11
+ doc = nlp(text)
12
+
13
+ pos_tokens = []
14
+
15
+ for token in doc:
16
+ if token.ent_type_ != "":
17
+ pos_tokens.append((token.text, token.ent_type_))
18
+ else:
19
+ pos_tokens.append((token.text, None))
20
+ return pos_tokens
21
+ elif model_type == "Roberta":
22
+ output = ner_pipeline(text)
23
+ print(output)
24
+ return {"text": text, "entities": [
25
+ {"word": entity["word"], "entity": entity["entity_group"], "start": entity['start'],
26
+ 'end': entity['end']}
27
+ for entity in output]}
28
+ elif model_type == "AWS Comprehend":
29
+ client = boto3.client('comprehend')
30
+ response = client.detect_entities(
31
+ Text=text, LanguageCode='en')
32
+ print(response)
33
+ return {"text": text, "entities": [{"word": entity["Text"], "entity": entity["Type"], "start": entity['BeginOffset'], 'end': entity['EndOffset']}
34
+ for entity in response["Entities"]]}
35
+
36
+
37
+ demo = gr.Interface(fn=greet, inputs=[gr.Radio(["Spacy", "Roberta", "AWS Comprehend"]), "text"],
38
+ outputs="highlight", title="Key Matcher",
39
+ description=f"Find the best match for a key from the list of ",)
40
+
41
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ spacy
2
+ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.5.0/en_core_web_sm-3.5.0-py3-none-any.whl
3
+ transformers
4
+ boto3