From 775313982f6c23f222039a905b3b191219f06c3c Mon Sep 17 00:00:00 2001
From: MaartenGr <maarten_grootendorst@hotmail.com>
Date: Thu, 7 Dec 2023 12:44:43 +0100
Subject: [PATCH] Fix extraction
---
keybert/llm/_openai.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/keybert/llm/_openai.py b/keybert/llm/_openai.py
index df665bdb..5c8c078c 100644
--- a/keybert/llm/_openai.py
+++ b/keybert/llm/_openai.py
@@ -178,7 +178,7 @@ def extract_keywords(self, documents: List[str], candidate_keywords: List[List[s
response = chat_completions_with_backoff(self.client, **kwargs)
else:
response = self.client.chat.completions.create(**kwargs)
- keywords = response["choices"][0]["message"]["content"].strip()
+ keywords = response.choices[0].message.content.strip()
# Use a non-chat model
else:
@@ -186,7 +186,7 @@ def extract_keywords(self, documents: List[str], candidate_keywords: List[List[s
response = completions_with_backoff(self.client, model=self.model, prompt=prompt, **self.generator_kwargs)
else:
response = self.client.completions.create(model=self.model, prompt=prompt, **self.generator_kwargs)
- keywords = response["choices"][0]["text"].strip()
+ keywords = response.choices[0].message.content.strip()
keywords = [keyword.strip() for keyword in keywords.split(",")]
all_keywords.append(keywords)