Repository: shreyashankar/gpt3-sandbox
File: api/gpt.py (Python)

Code before:
"""Creates the Example and GPT classes for a user to interface with the OpenAI API."""

import openai


def set_openai_key(key):
    """Sets OpenAI key."""
    openai.api_key = key

class Example():
    """Stores an input, output pair and formats it to prime the model."""

    def __init__(self, inp, out, input_prefix = "input: ", input_suffix = "\n", output_prefix = "output: ", output_suffix = "\n"):
        self.input = inp
        self.output = out
        self.input_prefix = input_prefix
        self.input_suffix = input_suffix
        self.output_prefix = output_prefix
        self.output_suffix = output_suffix

    def get_input(self):
        """Returns the input of the example."""
        return self.input

    def get_output(self):
        """Returns the intended output of the example."""
        return self.output

    def format(self):
        """Formats the input, output pair."""
        return self.input_prefix + f"{self.input}" + self.input_suffix + self.output_prefix + f"{self.output}\n"
        + self.output_suffix


class GPT:
    """The main class for a user to interface with the OpenAI API.
    A user can add examples and set parameters of the API request."""

    def __init__(self, engine='davinci',
                 temperature=0.5,
                 max_tokens=100,
                 input_prefix = "input: ",
                 input_suffix = "\n",
                 output_prefix = "output: ",
                 output_suffix = "\n",
                 append_input_suffix_and_output_prefix_to_query = False,
                 stop = "\ninput:"):
        self.examples = []
        self.engine = engine
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.input_prefix = input_prefix
        self.input_suffix = input_suffix
        self.output_prefix = output_prefix
        self.output_suffix = output_suffix
        self.append_input_suffix_and_output_prefix_to_query = append_input_suffix_and_output_prefix_to_query
        self.stop = stop

    def add_example(self, ex):
        """Adds an example to the object. Example must be an instance
        of the Example class."""
        assert isinstance(ex, Example), "Please create an Example object."
        self.examples.append(ex.format())

    def get_prime_text(self):
        """Formats all examples to prime the model."""
        return '\n'.join(self.examples) + '\n'

    def get_engine(self):
        """Returns the engine specified for the API."""
        return self.engine

    def get_temperature(self):
        """Returns the temperature specified for the API."""
        return self.temperature

    def get_max_tokens(self):
        """Returns the max tokens specified for the API."""
        return self.max_tokens

    def craft_query(self, prompt):
        """Creates the query for the API request."""
        q = self.get_prime_text() + self.input_prefix + prompt
        if self.append_input_suffix_and_output_prefix_to_query:
            q = q + self.input_suffix + self.output_prefix

        return q

    def submit_request(self, prompt):
        """Calls the OpenAI API with the specified parameters."""
        response = openai.Completion.create(engine=self.get_engine(),
                                            prompt=self.craft_query(prompt),
                                            max_tokens=self.get_max_tokens(),
                                            temperature=self.get_temperature(),
                                            top_p=1,
                                            n=1,
                                            stream=False,
                                            stop=self.stop)
        return response

    def get_top_reply(self, prompt):
        """Obtains the best result as returned by the API."""
        response = self.submit_request(prompt)
        return response['choices'][0]['text']


Review comment (suggestion):
In this class, you could include a method such as:

```
def format_example(self, ex):
        """Formats the input, output pair."""
        return self.input_prefix + ex.get_input() + self.input_suffix + self.output_prefix + ex.get_output() + self.output_suffix
```

Then the `add_example` method could be:

```
def add_example(self, ex):
        """Adds an example to the object. Example must be an instance
        of the Example class."""
        assert isinstance(ex, Example), "Please create an Example object."
        self.examples.append(self.format_example(ex))
```

Code after:

    def format_query(self, prime_text, prompt):
        """Creates the query for the API request."""
        query = prime_text + self.input_format.get_wrapped_text(prompt)
        if self.append_output_prefix_to_query:
            query = query + self.output_format.get_prefix()
        return query

class GPT:
    """The main class for a user to interface with the OpenAI API.
    A user can add examples and set parameters of the API request."""

    def __init__(self, engine='davinci',
                 temperature=0.5,
                 max_tokens=100,
                 query_format=QueryFormat()):
        self.examples = []
        self.engine = engine
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.query_format = query_format

        assert isinstance(query_format, QueryFormat), \
        "Please create a QueryFormat object."

    def add_example(self, ex):
        """Adds an example to the object. Example must be an instance
        of the Example class."""
        assert isinstance(ex, Example), "Please create an Example object."
        self.examples.append(self.query_format.format_example(ex))

    def get_prime_text(self):
        """Formats all examples to prime the model."""
        return ''.join(self.examples)

    def get_engine(self):
        """Returns the engine specified for the API."""
        return self.engine

    def get_temperature(self):
        """Returns the temperature specified for the API."""
        return self.temperature

    def get_max_tokens(self):
        """Returns the max tokens specified for the API."""
        return self.max_tokens

    def get_append_output_prefix_to_query(self):
        """Returns the suffix of this format."""
        return self.query_format.get_append_output_prefix_to_query()