Add multiple language models choices

- The user can select between 'text-davinci-003' and 'gpt-3.5-turbo' - Fix an issue with automatic #flashcards tag generation
crybot · Jul 5, 2023 · e1b5bb0 · e1b5bb0
1 parent 79c0c8f
commit e1b5bb0
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 22 deletions.
diff --git a/src/flashcards.ts b/src/flashcards.ts
@@ -20,36 +20,54 @@ function extractTextAfterFlashcards(text: string): string | null {
 }
 
 
-export async function generateFlashcards(text: string, apiKey: string): Promise<string> {
+export async function generateFlashcards(text: string, apiKey: string, model: string = 'text-davinci-003'): Promise<string> {
   const configuration = new Configuration({
     apiKey: apiKey,
   });
 
   const openai = new OpenAIApi(configuration);
 
   const cleanedText = text.replace(/<!--.*-->[\n]?/g, "");
-  const alreadyGenerated = extractTextAfterFlashcards(cleanedText);
-  const flashcardText = cleanedText.split("#flashcards")[0];
+  const flashcardText = cleanedText
 
-  const basePrompt = `I'll provide you with a note, identify which are the most important concepts within it and generate at most 5 flashcards in the format \"question :: answer\". Strictly use :: to separate a question from its answer.An example is \"What is chemical formula of water :: H2O\". Do not use any prefix text, start generating right away. Try to make them as atomic as possible, but still challenging and rich of information. Do not generate flashcards that have already been generated. DO NOT REPEAT OR REPHRASE ANY OF THE FOLLOWING flashcards already generated:${alreadyGenerated}`;
+  const basePrompt = `I'll provide you with a note. At the end of the note are some flashcards. Identify which are the most important concepts within the note and generate 3 new original flashcard in the format \"question :: answer\". Strictly use :: to separate a question from its answer. Separate flashcards with a single newline. An example is \"What is chemical formula of water :: H2O\". Do not use any prefix text, start generating right away. Try to make them as atomic as possible, but still challenging and rich of information. DO NOT REPEAT OR REPHRASE FLASHCARDS. Focus on important latex formulas and equations. Please typeset equations and math formulas correctly (that is using the \$ symbol)`;
   const additionalPrompt = "Additional information on the task: Focus primarily on formulas and equations. Do NOT always start the questions with What. Do not repeat questions. Do not rephrase questions already generated. You can also ask the user to describe something or detail a given concept. You can even write flashcards asking to fill a missing word or phrase.";
 
-  const prompt = `${basePrompt}\n${flashcardText}`;
+  let response = null;
+  if (model == 'gpt-3.5-turbo') {
+    response = await openai.createChatCompletion({
+          model: model,
+          temperature: 0.7,
+          max_tokens: 300,
+          frequency_penalty: 0,
+          presence_penalty: 0,
+          top_p: 1.0,
+          messages: [{role: "system", content: basePrompt}, {role: "user", content: flashcardText}],
+      }, { timeout: 60000 });
+  }
 
-  const response = await openai.createCompletion({
-    model: "text-davinci-003",
-    prompt: prompt,
-    temperature: 0.3,
-    max_tokens: 500,
-    top_p: 1.0,
-    frequency_penalty: 1.3,
-    presence_penalty: 0.6,
-  });
+  else if(model == 'text-davinci-003') {
+    const prompt = `${basePrompt}\n${flashcardText}`;
+
+    response = await openai.createCompletion({
+      model: model,
+      prompt: prompt,
+      temperature: 0.7,
+      max_tokens: 300,
+      top_p: 1.0,
+      frequency_penalty: 0,
+      presence_penalty: 0,
+    });
+  }
 
-  if (response.data.choices && response.data.choices.length > 0) {
-    return response.data.choices[0].text;
-  } else {
-    console.log(response)
-    throw new OpenAIError("No response received from OpenAI API");
+
+  if (model == 'text-davinci-003') {
+    return response.data.choices[0].text.trim();
+  }
+  else if (model == 'gpt-3.5-turbo') {
+    return response.data.choices[0].message.content.trim();
   }
+  throw new OpenAIError("No response received from OpenAI API");
+  console.log(response)
+
 }
diff --git a/src/main.ts b/src/main.ts
@@ -3,10 +3,12 @@ import { generateFlashcards } from "./flashcards";
 
 interface FlashcardsSettings {
   apiKey: string;
+  model: string;
 }
 
 const DEFAULT_SETTINGS: FlashcardsSettings = {
   apiKey: "",
+  model: "text-davinci-003"
 };
 
 export default class FlashcardsLLMPlugin extends Plugin {
@@ -29,6 +31,7 @@ export default class FlashcardsLLMPlugin extends Plugin {
 
   async onGenerateFlashcards(editor: Editor, view: MarkdownView) {
     const apiKey = this.settings.apiKey;
+    const model = this.settings.model;
     if (!apiKey) {
       new Notice("API key is not set in plugin settings");
       return;
@@ -40,7 +43,7 @@ export default class FlashcardsLLMPlugin extends Plugin {
     const hasHeader = headerRegex.test(currentText);
 
     // Check if the #flashcards tag is already present
-    const tagRegex = /\n#flashcards\n/;
+    const tagRegex = /\n#flashcards.*\n/;
     const hasTag = tagRegex.test(currentText);
 
     let updatedText = currentText;
@@ -57,8 +60,8 @@ export default class FlashcardsLLMPlugin extends Plugin {
 
     new Notice("Generating flashcards...");
     try {
-      const generatedCards = (await generateFlashcards(updatedText, apiKey)).split("\n");
-      editor.setValue(updatedText + "\n" + generatedCards.join('\n\n'))
+      const generatedCards = (await generateFlashcards(updatedText, apiKey, model)).split("\n");
+      editor.setValue(updatedText + "\n\n" + generatedCards.map(s => s.trim()).join('\n\n'))
 
       const newPosition: EditorPosition = {
         line: editor.lastLine()
@@ -110,5 +113,19 @@ class FlashcardsSettingsTab extends PluginSettingTab {
         await this.plugin.saveSettings();
       })
     );
+
+    new Setting(containerEl)
+    .setName("Model")
+    .setDesc("Which language model to use")
+    .addDropdown((dropdown) =>
+      dropdown
+      .addOption("text-davinci-003", "text-davinci-003")
+      .addOption("gpt-3.5-turbo", "gpt-3.5-turbo")
+      .setValue(this.plugin.settings.model)
+      .onChange(async (value) => {
+        this.plugin.settings.model = value;
+        await this.plugin.saveSettings();
+      })
+    );
   }
 }