131 lines
3.6 KiB
Python
131 lines
3.6 KiB
Python
import openai
|
|
import os
|
|
import json
|
|
import genanki
|
|
from logging_config import setup_logging
|
|
from constants import API_KEY_ENV, CHAT_MODEL
|
|
|
|
|
|
setup_logging()
|
|
|
|
|
|
API_KEY = os.environ.get(API_KEY_ENV)
|
|
if not API_KEY:
|
|
raise ValueError("Please set the OPENAI_API_KEY environment variable.")
|
|
|
|
openai.api_key = API_KEY
|
|
|
|
# Given prompt template
|
|
PROMPT_TEMPLATE = """
|
|
Please come up with a title for the deck and a set of 10 index cards for memorization,
|
|
including a title, front, and back for each card. The index cards should completely
|
|
capture the main points and themes of the text. In addition, they should contain any
|
|
numbers or data that humans might find difficult to remember. The goal of the index
|
|
card set is that one who memorizes it can provide a summary of the text to someone
|
|
else, conveying the main points and themes.
|
|
|
|
You will provide the deck title, and the titles, questions, and answers for each card
|
|
in a structured format as follows:
|
|
```
|
|
Deck Title: Title of the Deck
|
|
Cards:
|
|
- Title: Card Title 1
|
|
Front: What is the capital of New York?
|
|
Back: Albany
|
|
- Title: Card Title 2
|
|
Front: Where in the world is Carmen San Diego?
|
|
Back: Nobody knows
|
|
```
|
|
|
|
{content}
|
|
"""
|
|
|
|
|
|
def prompt_for_card_content(text_content):
|
|
# Prepare the prompt
|
|
prompt = PROMPT_TEMPLATE.format(content=text_content)
|
|
|
|
# Get completion from the OpenAI ChatGPT API
|
|
response = openai.ChatCompletion.create(
|
|
model=CHAT_MODEL,
|
|
messages=[
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
temperature=0,
|
|
)
|
|
|
|
# Extract content from response and save to a new file
|
|
return response.choices[0]['message']['content']
|
|
|
|
|
|
def response_to_json(response_text):
|
|
lines = [line.strip() for line in response_text.split("\n") if line.strip()]
|
|
|
|
deck_title = None
|
|
cards = []
|
|
current_card = {}
|
|
|
|
for line in lines:
|
|
if "Deck Title:" in line and not deck_title:
|
|
deck_title = line.split("Deck Title:", 1)[1].strip()
|
|
elif "Title:" in line:
|
|
if current_card: # If there's a card being processed, add it to cards
|
|
cards.append(current_card)
|
|
current_card = {}
|
|
current_card["Title"] = line.split("Title:", 1)[1].strip()
|
|
elif "Front:" in line:
|
|
current_card["Question"] = line.split("Front:", 1)[1].strip()
|
|
elif "Back:" in line:
|
|
current_card["Answer"] = line.split("Back:", 1)[1].strip()
|
|
|
|
if current_card: # Add the last card if it exists
|
|
cards.append(current_card)
|
|
|
|
return {
|
|
"DeckTitle": deck_title,
|
|
"Cards": cards
|
|
}
|
|
|
|
|
|
# Create a new model for our cards. This is necessary for genanki.
|
|
MY_MODEL = genanki.Model(
|
|
1607372319,
|
|
"Simple Model",
|
|
fields=[
|
|
{"name": "Title"},
|
|
{"name": "Question"},
|
|
{"name": "Answer"},
|
|
],
|
|
templates=[
|
|
{
|
|
"name": "{{Title}}",
|
|
"qfmt": "{{Question}}",
|
|
"afmt": "{{FrontSide}}<hr id='answer'>{{Answer}}",
|
|
},
|
|
])
|
|
|
|
def json_file_to_package(json_path):
|
|
with open(json_path, 'r', encoding='utf-8') as f:
|
|
json_data = json.load(f)
|
|
package = to_package(json_data)
|
|
|
|
return package
|
|
|
|
def to_package(deck_json):
|
|
deck_title = deck_json["DeckTitle"]
|
|
deck = genanki.Deck(1607372319, deck_title)
|
|
|
|
for card_json in deck_json["Cards"]:
|
|
title = card_json["Title"]
|
|
question = card_json["Question"]
|
|
answer = card_json["Answer"]
|
|
|
|
note = genanki.Note(
|
|
model=MY_MODEL,
|
|
fields=[title, question, answer]
|
|
)
|
|
|
|
deck.add_note(note)
|
|
|
|
return genanki.Package(deck)
|