use new 3.5 turbo chat model.

1 files changed, 17 insertions, 18 deletions

diff --git a/setup.py b/setup.py
index 9efb57c..7bbfd6d 100644
--- a/setup.py
+++ b/setup.py
@@ -146,21 +146,20 @@ def generate_summary(prompt):
     if (len(enc.encode(prompt)) > 3000):
         return "too long to summarize."
 
-    prompt = prompt + '\nSummarize the above code: '
-
- # response = openai.ChatCompletion.create(
- #   model="gpt-3.5-turbo",
- #   messages=[{"role": "user", "content": prompt}],
- #   temperature=0.7,
- #   max_tokens=1024,
- #   top_p=1.0,
- #   frequency_penalty=0.0,
- #   presence_penalty=0.0,
- #   stop=["\"\"\""]
- # )
-
-    #return response["choices"][0]["message"]["content"]
-    return 'herro. this is a test summary'
+    prompt = prompt + '\nSummarize the above code (be succinct): '
+
+    response = openai.ChatCompletion.create(
+      model="gpt-3.5-turbo",
+      messages=[{"role": "user", "content": prompt}],
+      temperature=0.7,
+      max_tokens=300,
+      top_p=1.0,
+      frequency_penalty=0.0,
+      presence_penalty=0.0,
+      stop=["\"\"\""]
+    )
+
+    return response["choices"][0]["message"]["content"]
 
 # create blob. the blob just contains the file path and the source code.
 def blobify(pandaSeries):
@@ -234,9 +233,9 @@ def setup(
 
     print('generating embeddings')
     embedding_model = "text-embedding-ada-002"
-    #code_df["embedding_summary"] = code_df.summary.apply(
-    #        [lambda x: get_embedding(x, engine=embedding_model)]
-    #        )
+    code_df["embedding_summary"] = code_df.summary.apply(
+            [lambda x: get_embedding(x, engine=embedding_model)]
+            )
     print('done with embeddings')
     code_df.to_csv(output_csv_filepath)