対象フォルダから特定文字列のみを抽出するスクリプト

はじめに

自然言語領域で使ったスクリプトの一部です。

用途としては、Wikipediaから抽出された膨大なテキストデータ(9GB...)から特定文字列を含む行のみを抽出して、それをAImodelの学習データに組み込みたかったため、突貫でつくりました。

表示するコード

print_words.py

import os
import re

search_dir = "/home/ono/test"
search_pattern = ""

file_name_list = os.listdir(search_dir)
for file_name in file_name_list:
    f = open(os.path.join(search_dir, file_name))
    line = f.readline()
    line_number = 1
    while line:
        m = re.search(search_pattern, line)
        if m:
            print("Pattern found: file:%s, line:%d, data:%s" % (file_name, line_number, line))
        line = f.readline()
        line_number = line_number + 1
    f.close()

少しのメモ書き

os.listdir

ディレクトリとファイルの一覧を取得。
その後for文で一文ずつファイル名のみ取り出す。

os.path.join(x,y)

ディレクトリ+特定ファイル名で走査ファイルの決定を行う。

line_number = 1 の理由

number は未定義の場合、0始まりなので、表示用計算ならば1を置いた。

re.search(x, y)

文章がパターンにマッチするか検索している部分

今振り返ると改善点。

• 検索文はsys.args[]などで引数指定にすることも可能
• re.search()よりもre.compile()にすると実行速度が向上する
• open()よりもwith open()のほうがコードが短くなる

抽出するコード

get_words.py

import os
from glob import glob

input_dir = "/home/ono/AA"
output_file = "/home/ono/extract.txt"
featuredwords = "dicTechterm_180411.txt"

def main():

    hako = []

    for path in glob(os.path.join(input_dir,'wiki_*')):
        with open(path,"r",encoding="utf-8") as file:
            for contents in iter_docs(file):
                if contents != "":
                    hako.append(contents)
                with open(output_file, "w") as out:
                    out.write('\n'.join(hako))

    print(hako)
    
def iter_docs(file):

    contents = []

    with open(featuredwords, "r",encoding="utf-8") as term:
        #特徴語を格納する
        query = []
        for text in term:
            query.append(text.strip('\n'))

        #格納した特徴語を検索キーとして、
        for w in query:
            breakflg = 0 #初期化フラグ
            for line in file:
                if breakflg == 1:
                    continue
                if line.startswith('<doc '):
                    pass
                elif line.startswith('</doc>'):
                    pass
                else:
                    if line.find(w) > -1:
                        contents.append(line.strip("\n")) 
                        breakflg = 1
                        break
    return contents
                
if __name__ == '__main__':
    main()