Merge pull request #16 from shreyashankar/buyumaz/move_api_call_to_GP…

…T_class moved api logic into GPT class
shreyashankar · Jul 19, 2020 · db62d15 · db62d15
2 parents 3577cc5 + 9ca9ed7
commit db62d15
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 14 deletions.
diff --git a/api/GPT.py b/api/GPT.py
@@ -1,3 +1,5 @@
+import openai
+
 class Example():
     def __init__(self, inp, out):
         self.input = inp
@@ -16,6 +18,9 @@ def __init__(self, engine='davinci',
         self.temperature=temperature
         self.max_tokens=max_tokens
 
+    def set_openai_key(self, key):
+        openai.api_key=key
+
     def add_example(self, ex):
         assert isinstance(ex, Example), "Please create an Example object."
         self.examples.append(ex.format())
@@ -34,3 +39,14 @@ def get_max_tokens(self):
 
     def craft_query(self, prompt):
         return self.get_prime_text() + "input: " + prompt + "\n"
+
+    def submit_request(self, prompt):
+        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="\ninput:")
+        return response 
diff --git a/api/demo_web_app.py b/api/demo_web_app.py
@@ -4,23 +4,11 @@
 
 from .UIConfig import UIConfig
 
-def submit_openai_request(gpt, prompt):
-    response = openai.Completion.create(engine=gpt.get_engine(),
-                                        prompt=gpt.craft_query(prompt),
-                                        max_tokens=gpt.get_max_tokens(),
-                                        temperature=gpt.get_temperature(),
-                                        top_p=1,
-                                        n=1,
-                                        stream=False,
-                                        stop="\ninput:") 
-    return response
-
-
 def demo_web_app(gpt, config=UIConfig()):
     app = Flask(__name__)
     app.config.from_envvar('LATEX_TRANSLATOR_CONFIG')
     CORS(app)
-    openai.api_key = app.config['OPENAI_KEY']
+    gpt.set_openai_key(app.config['OPENAI_KEY'])
 
     @app.route("/params", methods=['GET'])
     def get_params():
@@ -30,7 +18,7 @@ def get_params():
     @app.route("/translate", methods=['GET', 'POST'])
     def translate():
         prompt = request.json['prompt']
-        response = submit_openai_request(gpt, prompt)
+        response = gpt.submit_request(prompt) 
         return {'text': response['choices'][0]['text'][7:]}
 
     app.run()