-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdata.py
More file actions
121 lines (104 loc) 路 4.1 KB
/
Copy pathdata.py
File metadata and controls
121 lines (104 loc) 路 4.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import os
import re
import codecs
import logging
import csv
def log_lines(file, n=10):
"""
Preview n lines from file
:param file: str
:param n: int
:return:
"""
logging.info("Printing some lines from {}:".format(file))
with open(file, 'rb') as datafile:
lines = datafile.readlines()
for line in lines[:n]:
logging.info("\t {}".format(line))
def load_lines(filename, fields):
"""
Splits each line of the file into a dictionary of fields.
Load a line formatted file into a dictionary of line objects.
:param filename: str
:param fields: list[str]
:return: dict{lineID: line_obj}
"""
lines = {}
with open(filename, 'r', encoding='iso-8859-1') as f:
for line in f:
values = line.split(" +++$+++ ")
# Extract fields
line_obj = {}
for i, field in enumerate(fields):
line_obj[field] = values[i]
lines[line_obj['lineID']] = line_obj
return lines
def load_conversations(filename, lines, fields):
"""
Group fields of lines from 'load_lines' into conversations based on *movie_conversations.txt*
Load a conversation formatted file into a list of conversation objects.
:param filename: str
:param lines:
:param fields: list[str]
:return: list[dict]
"""
conversations = []
with open(filename, 'r', encoding='iso-8859-1') as f:
for line in f:
values = line.split(" +++$+++ ")
# Extract fields
conv_obj = {}
for i, field in enumerate(fields):
conv_obj[field] = values[i]
# Convert string to list (conv_obj["utteranceIDs"] == "['L598485', 'L598486', ...]")
utterance_id_pattern = re.compile('L[0-9]+')
line_ids = utterance_id_pattern.findall(conv_obj["utteranceIDs"])
# Reassemble lines
conv_obj["lines"] = []
for line_id in line_ids:
conv_obj["lines"].append(lines[line_id])
conversations.append(conv_obj)
return conversations
def extract_sentence_pairs(conversations):
"""
Extract paris of sentences from conversations. Return question answer pairs where the target line
is the line after the input line in a conversation.
:param conversations: list[conv_obj]
:return: list[list[str,str]]
"""
qa_pairs = []
for conversation in conversations:
# Iterate over all the lines of the conversation
for i in range(len(conversation["lines"]) - 1): # We ignore the last line (no answer for it)
input_line = conversation["lines"][i]["text"].strip()
target_line = conversation["lines"][i + 1]["text"].strip()
# Filter wrong samples (if one of the lists is empty)
if input_line and target_line:
qa_pairs.append([input_line, target_line])
return qa_pairs
def create_formatted_file(config):
"""
Create a formatted .csv file of the data.
:return: None
"""
datafile = config["datafile"]
delimiter = "~"
# Unescape the delimiter
delimiter = str(codecs.decode(delimiter, "unicode_escape"))
# Initialize field ids
MOVIE_LINES_FIELDS = ["lineID", "characterID", "movieID", "character", "text"]
MOVIE_CONVERSATIONS_FIELDS = ["character1ID", "character2ID", "movieID", "utteranceIDs"]
# Load lines and process conversations
logging.info("Processing corpus")
lines = load_lines(os.path.join(config["corpus"], "movie_lines.txt"), MOVIE_LINES_FIELDS)
logging.info("Loading conversations...")
conversations = load_conversations(os.path.join(config["corpus"], "movie_conversations.txt"), lines,
MOVIE_CONVERSATIONS_FIELDS)
# Write new csv file
logging.info("Writing newly formatted file...")
with open(datafile, 'w', encoding='utf-8') as outfile:
writer = csv.writer(outfile, delimiter=delimiter, lineterminator='\n')
for pair in extract_sentence_pairs(conversations):
writer.writerow(pair)
# Print some samples of lines from the new file
log_lines(datafile)