Spaces:
Running
on
Zero
Running
on
Zero
Commit
·
123d671
1
Parent(s):
81fbf66
Increase duration
Browse filesSigned-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>
app.py
CHANGED
|
@@ -36,7 +36,7 @@ def get_country(country_name: str):
|
|
| 36 |
return None
|
| 37 |
|
| 38 |
|
| 39 |
-
@spaces.GPU(duration=
|
| 40 |
def predict_entities(model, query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False):
|
| 41 |
start = datetime.now()
|
| 42 |
|
|
@@ -53,7 +53,8 @@ def predict_entities(model, query: str, labels: Union[str, list], threshold: flo
|
|
| 53 |
def parse_query(query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False, model_name: str = None) -> Dict[str, Union[str, list]]:
|
| 54 |
|
| 55 |
entities = []
|
| 56 |
-
|
|
|
|
| 57 |
|
| 58 |
for entity in _entities:
|
| 59 |
if entity["label"] == "country":
|
|
|
|
| 36 |
return None
|
| 37 |
|
| 38 |
|
| 39 |
+
@spaces.GPU(duration=10)
|
| 40 |
def predict_entities(model, query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False):
|
| 41 |
start = datetime.now()
|
| 42 |
|
|
|
|
| 53 |
def parse_query(query: str, labels: Union[str, list], threshold: float = 0.3, nested_ner: bool = False, model_name: str = None) -> Dict[str, Union[str, list]]:
|
| 54 |
|
| 55 |
entities = []
|
| 56 |
+
model = get_model(model_name)
|
| 57 |
+
_entities = predict_entities(model=model, query=query, labels=labels, threshold=threshold, nested_ner=nested_ner)
|
| 58 |
|
| 59 |
for entity in _entities:
|
| 60 |
if entity["label"] == "country":
|