Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Nov 14, 2024
1 parent 099086f commit 5e19355
Showing 1 changed file with 8 additions and 7 deletions.
15 changes: 8 additions & 7 deletions torch_geometric/nn/nlp/txt2kg.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,18 +54,19 @@ def chunks_to_triples_strs(self, txt_batch: List[str]) -> List[str]:
self.model = LLM(LM_name, num_params=14).eval()
self.initd_LM = True
out_strs = self.model.inference(
question=[txt + '\n' + self.system_prompt for txt in txt_batch],
max_tokens=self.chunk_size)
question=[
txt + '\n' + self.system_prompt for txt in txt_batch
], max_tokens=self.chunk_size)
else:
messages = []
for txt in txt_batch:
messages.append({
"role":
"user",
"content":
txt + '\n' + self.system_prompt})
"role": "user",
"content": txt + '\n' + self.system_prompt
})
completion = self.client.chat.completions.create(
model=self.model, messages=messages, temperature=0, top_p=1, max_tokens=1024, stream=True)
model=self.model, messages=messages, temperature=0, top_p=1,
max_tokens=1024, stream=True)
out_str = ""
for chunk in completion:
if chunk.choices[0].delta.content is not None:
Expand Down

0 comments on commit 5e19355

Please sign in to comment.