textprocessor.py

import os
import json
from typing import Dict

import openai

context_prompt = """
You are a creator of illustrated books building a series of scenes for your book.
Your boss asked you to write a summary and illustrations of this text:
$TRANSCRIPTION
You have to write the summary using a maximum of 7 scenes in a JSON object with these keys:
"Scene": The number of the scene.
"Summary": Spanish string with a summary of the scene. It should be in Spanish, and it should be less than 30 words. Readers should understand it without looking at the illustration.
"Illustration": English string with a detailed English description of an illustration for this scene. It must be written in English and in less than 20 words. It should include many details and an artistic style for the image that matches the text.
Just answer with the JSON object:
"""

openai.api_key = os.getenv("SECRET_KEY_OPENAI")

class TextProcessor:
    def __init__(self,
                 model: str = "text-davinci-003",
                 temperature: float = 0.7,
                 max_tokens: int = 1240,
                 top_p: int = 1,
                 frequency_penalty: int = 0,
                 presence_penalty: int = 0) -> None:
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.top_p = top_p
        self.frequency_penalty = frequency_penalty
        self.presence_penalty = presence_penalty
    
    def get_json_scenes(self,
                        prompt: str) -> Dict:
        gpt_prompt = context_prompt.replace("$TRANSCRIPTION", prompt)
        print("gpt_prompt", gpt_prompt)
        response = openai.Completion.create(
            model=self.model,
            prompt=gpt_prompt,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            top_p=self.top_p,
            frequency_penalty=self.frequency_penalty,
            presence_penalty=self.presence_penalty
        )
        print("GPT response", response)
        scenes = json.loads(response["choices"][0]["text"])
        print("scenes", scenes)
        if (type(scenes) == list):
            scenes = {i: d for i, d in enumerate(scenes)}
        return scenes