Skip to content

Commit 32f9b99

Browse files
Merge pull request stanford-oval#189 from xiaobeicn/main
fix encoder documentation; add costorm-runner instance dumping in the example script
2 parents 735f547 + c40cbfb commit 32f9b99

File tree

3 files changed

+7
-27
lines changed

3 files changed

+7
-27
lines changed

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -255,7 +255,7 @@ python examples/storm_examples/run_storm_wiki_gpt.py \
255255
256256
To run Co-STORM with `gpt` family models with default configurations,
257257
258-
1. Add `BING_SEARCH_API_KEY="xxx"`to `secrets.toml`
258+
1. Add `BING_SEARCH_API_KEY="xxx"` and `ENCODER_API_TYPE="xxx"` to `secrets.toml`
259259
2. Run the following command
260260
261261
```bash

examples/costorm_examples/run_costorm_gpt.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,11 @@ def main(args):
143143
with open(os.path.join(args.output_dir, "report.md"), "w") as f:
144144
f.write(article)
145145

146+
# Save instance dump
147+
instance_copy = costorm_runner.to_dict()
148+
with open(os.path.join(args.output_dir, "instance_dump.json"), "w") as f:
149+
json.dump(instance_copy, f, indent=2)
150+
146151
# Save logging
147152
log_dump = costorm_runner.dump_logging_and_reset()
148153
with open(os.path.join(args.output_dir, "log.json"), "w") as f:

knowledge_storm/encoder.py

Lines changed: 1 addition & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -7,38 +7,13 @@
77

88

99
class EmbeddingModel:
10-
def __init__():
10+
def __init__(self):
1111
pass
1212

1313
def get_embedding(self, text: str) -> Tuple[np.ndarray, int]:
1414
raise Exception("Not implemented")
1515

1616

17-
class OpenAIEmbeddingModel(EmbeddingModel):
18-
def __init__(self, model: str = "text-embedding-3-small", api_key: str = None):
19-
if not api_key:
20-
self.api_key = os.getenv("OPENAI_API_KEY")
21-
22-
self.url = "https://api.openai.com/v1/embeddings"
23-
self.headers = {
24-
"Content-Type": "application/json",
25-
"Authorization": f"Bearer {self.api_key}",
26-
}
27-
self.model = model
28-
29-
def get_embedding(self, text: str) -> Tuple[np.ndarray, int]:
30-
data = {"input": text, "model": "text-embedding-3-small"}
31-
32-
response = requests.post(self.url, headers=self.headers, json=data)
33-
if response.status_code == 200:
34-
data = response.json()
35-
embedding = np.array(data["data"][0]["embedding"])
36-
token = data["usage"]["prompt_tokens"]
37-
return embedding, token
38-
else:
39-
response.raise_for_status()
40-
41-
4217
class OpenAIEmbeddingModel(EmbeddingModel):
4318
def __init__(self, model: str = "text-embedding-3-small", api_key: str = None):
4419
if not api_key:

0 commit comments

Comments
 (0)