forked from noahshinn/reflexion
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
41 lines (35 loc) · 1.21 KB
/
utils.py
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
import os
import gzip
import json
import openai
import jsonlines
from typing import List
openai.api_key = os.getenv("OPENAI_API_KEY")
def read_jsonl(path: str) -> List[dict]:
if not os.path.exists(path):
raise FileNotFoundError(f"File `{path}` does not exist.")
elif not path.endswith(".jsonl"):
raise ValueError(f"File `{path}` is not a jsonl file.")
items = []
with jsonlines.open(path) as reader:
for item in reader:
items += [item]
return items
def write_jsonl(path: str, data: List[dict], append: bool = False):
if not path.endswith(".jsonl"):
raise ValueError(f"File `{path}` is not a jsonl file.")
with jsonlines.open(path, mode='a' if append else 'w') as writer:
for item in data:
writer.write(item)
def read_jsonl_gz(path: str) -> List[dict]:
if not path.endswith(".jsonl.gz"):
raise ValueError(f"File `{path}` is not a jsonl.gz file.")
with gzip.open(path, "rt") as f:
data = [json.loads(line) for line in f]
return data
def parse_body(text):
lines = text.split('\n')
for i in range(len(lines)-1, -1, -1):
if 'return' in lines[i]:
return '\n'.join(lines[:i+1])
return text