Update dataset_extraction.py
Browse files- dataset_extraction.py +3 -1
dataset_extraction.py
CHANGED
|
@@ -1,5 +1,6 @@
|
|
| 1 |
from typing import List, Tuple
|
| 2 |
import torch
|
|
|
|
| 3 |
from SciAssist import DatasetExtraction
|
| 4 |
|
| 5 |
device = "gpu" if torch.cuda.is_available() else "cpu"
|
|
@@ -7,7 +8,8 @@ de_pipeline = DatasetExtraction(os_name="nt")
|
|
| 7 |
|
| 8 |
|
| 9 |
def de_for_str(input):
|
| 10 |
-
|
|
|
|
| 11 |
|
| 12 |
# output = []
|
| 13 |
# for res in results["dataset_mentions"]:
|
|
|
|
| 1 |
from typing import List, Tuple
|
| 2 |
import torch
|
| 3 |
+
import nltk
|
| 4 |
from SciAssist import DatasetExtraction
|
| 5 |
|
| 6 |
device = "gpu" if torch.cuda.is_available() else "cpu"
|
|
|
|
| 8 |
|
| 9 |
|
| 10 |
def de_for_str(input):
|
| 11 |
+
list_input = nltk.sent_tokenize(input)
|
| 12 |
+
results = de_pipeline.extract(list_input, type="str", save_results=False)
|
| 13 |
|
| 14 |
# output = []
|
| 15 |
# for res in results["dataset_mentions"]:
|