wehub-resource-sync
37 lines
1.4 KiB
Python
37 lines
1.4 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2021-05-20 16:25
|
|
from typing import Union, List, Callable
|
|
|
|
from hanlp.common.dataset import TransformableDataset
|
|
from hanlp.utils.io_util import read_cells
|
|
|
|
STS_B_TRAIN = 'http://ixa2.si.ehu.es/stswiki/images/4/48/Stsbenchmark.tar.gz#sts-train.csv'
|
|
STS_B_DEV = 'http://ixa2.si.ehu.es/stswiki/images/4/48/Stsbenchmark.tar.gz#sts-dev.csv'
|
|
STS_B_TEST = 'http://ixa2.si.ehu.es/stswiki/images/4/48/Stsbenchmark.tar.gz#sts-test.csv'
|
|
|
|
|
|
class SemanticTextualSimilarityDataset(TransformableDataset):
|
|
def __init__(self,
|
|
data: Union[str, List],
|
|
sent_a_col,
|
|
sent_b_col,
|
|
similarity_col,
|
|
delimiter='auto',
|
|
transform: Union[Callable, List] = None,
|
|
cache=None,
|
|
generate_idx=None) -> None:
|
|
self.delimiter = delimiter
|
|
self.similarity_col = similarity_col
|
|
self.sent_b_col = sent_b_col
|
|
self.sent_a_col = sent_a_col
|
|
super().__init__(data, transform, cache, generate_idx)
|
|
|
|
def load_file(self, filepath: str):
|
|
for i, cells in enumerate(read_cells(filepath, strip=True, delimiter=self.delimiter)):
|
|
yield {
|
|
'sent_a': cells[self.sent_a_col],
|
|
'sent_b': cells[self.sent_b_col],
|
|
'similarity': float(cells[self.similarity_col])
|
|
}
|