Spaces:
Runtime error
Runtime error
| # | |
| # Pyserini: Reproducible IR research with sparse and dense representations | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| import json | |
| import os | |
| import unittest | |
| from shutil import rmtree | |
| class TestTokenizeJson(unittest.TestCase): | |
| def test_bert_single_file(self): | |
| inj = 'test_bert_single_file.json' | |
| outj = 'out_test_bert_single_file.json' | |
| f = open(inj, 'w') | |
| f.write('{"id": "doc1","contents": "I have a new gpu!"}\n{"id": "doc2","contents": "I do have an old gpu!"}') | |
| f.close() | |
| if(os.getcwd().endswith('tests')): | |
| os.system(f'python ../pyserini/tokenize_json_collection.py --input {inj} --output {outj}') | |
| else: | |
| os.system(f'python pyserini/tokenize_json_collection.py --input {inj} --output {outj}') | |
| with open(outj, 'r') as ret: | |
| for i, line in enumerate(ret): | |
| contents = json.loads(line)['contents'] | |
| if (i == 0): | |
| self.assertEqual('i have a new gp ##u !', contents) | |
| else: | |
| self.assertEqual('i do have an old gp ##u !', contents) | |
| ret.close() | |
| os.remove(inj) | |
| os.remove(outj) | |
| def test_bert_dir(self): | |
| indir = './test_tokenize_json' | |
| outdir = './test_out_tokenize_json' | |
| if(os.path.isdir(indir)): | |
| rmtree(indir) | |
| os.mkdir(indir) | |
| f1 = open(indir+'/doc00.json', 'w') | |
| f1.write('{"id": "doc1","contents": "I have a new gpu!"}\n{"id": "doc2","contents": "I do have an old gpu!"}') | |
| f1.close() | |
| f2 = open(indir+'/doc01.json', 'w') | |
| f2.write('{"id": "doc1","contents": "A new gpu!"}\n{"id": "doc2","contents": "An old gpu!"}') | |
| f2.close() | |
| if (os.getcwd().endswith('tests')): | |
| os.system(f'python ../pyserini/tokenize_json_collection.py --input {indir} --output {outdir}') | |
| else: | |
| os.system(f'python pyserini/tokenize_json_collection.py --input {indir} --output {outdir}') | |
| with open(outdir+'/docs00.json', 'r') as ret: | |
| for i, line in enumerate(ret): | |
| contents = json.loads(line)['contents'] | |
| if (i == 0): | |
| self.assertEqual('i have a new gp ##u !', contents) | |
| else: | |
| self.assertEqual('i do have an old gp ##u !', contents) | |
| with open(outdir+'/docs01.json', 'r') as ret: | |
| for i, line in enumerate(ret): | |
| contents = json.loads(line)['contents'] | |
| if (i == 0): | |
| self.assertEqual('a new gp ##u !', contents) | |
| else: | |
| self.assertEqual('an old gp ##u !', contents) | |
| rmtree(outdir) | |
| rmtree(indir) | |
| if __name__ == '__main__': | |
| unittest.main() | |