[bugfix] Bugfix

re-search · May 17, 2019 · 558d4a0 · 558d4a0
1 parent 458e52c
commit 558d4a0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/docproduct/train_embedding_to_gpt2_data.py b/docproduct/train_embedding_to_gpt2_data.py
@@ -36,8 +36,8 @@ def train_embedding_to_gpt2_data(
  question_index.add(question_bert)
 
  os.makedirs(os.path.dirname(output_path), exist_ok=True)
-
- output = open(output_path, "w")
+ csv_path = output_path + '.csv'
+ output = open(csv_path, "w")
  writer = csv.writer(output)
 
  firstrow = ['question', 'answer']
@@ -77,7 +77,7 @@ def topKforGPT2(start_ind, end_ind, topk):
  output.close()
 
  # ugly fix
- pd.read_csv(output_path, lineterminator='\n').to_parquet(
+ pd.read_csv(csv_path, lineterminator='\n').to_parquet(
  output_path, index=False)