Use hosted inference when on CPU

app.py

@@ -1,6 +1,8 @@
 import string
 import re
-from typing import List
+import json
+import requests
+from typing import List, Optional
 
 import torch
 from transformers import EncoderDecoderModel, BertTokenizerFast
@@ -60,11 +62,45 @@ class Predictor:
         return titles
 
 
+class HostedInference:
+
+    def __init__(self, model: str, num_titles: int, api_key: Optional[str] = None) -> None:
+        super().__init__()
+        self.model = model
+        self.num_titles = num_titles
+        self.api_key = api_key
+
+    def __call__(self, abstract: str, temperature: float) -> List[str]:
+        temperature = max(1.0, float(temperature))
+        data = json.dumps({
+            'inputs' : abstract,
+            'parameters' : {
+                'do_sample': (temperature > 1),
+                'num_beams': 10,
+                'temperature': temperature,
+                'top_k': 50,
+                'no_repeat_ngram_size': 2,
+                'num_return_sequences': self.num_titles,
+            },
+            'options' : { 'use_cache' : False, 'wait_for_model' : True }
+        })
+        api_url = "https://api-inference.huggingface.co/models/" + self.model
+        headers = { "Authorization": f"Bearer {self.api_key}" } if self.api_key is not None else {}
+        response = requests.request("POST", api_url, headers=headers, data=data)
+        response = json.loads(response.content.decode("utf-8"))
+
+        if isinstance(response, dict) and ('error' in response):
+            raise RuntimeError(response['error'])
+
+        titles = [str2title(title['summary_text']) for title in response]
+        return titles
+
+
 def create_gradio_ui(predictor):
 
     inputs = [
         gr.Textbox(label="Paper Abstract", lines=10),
-        gr.Slider(label="Creativity", minimum=1.0, maximum=2.0, step=0.1, value=1.5),
+        gr.Slider(label="Creativity", minimum=1.0, maximum=2.5, step=0.1, value=1.5),
     ]
     outputs = ["text"] * predictor.num_titles
 
@@ -86,13 +122,17 @@ def create_gradio_ui(predictor):
 
 
 if __name__ == '__main__':
-    print('Loading model...')
     model_path = "Callidior/bert2bert-base-arxiv-titlegen"
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    tokenizer = BertTokenizerFast.from_pretrained(model_path)
-    model = EncoderDecoderModel.from_pretrained(model_path).to(device)
-    print(f'Ready - running on {device}.')
+
+    if torch.cuda.is_available():
+        print('Loading model...')
+        tokenizer = BertTokenizerFast.from_pretrained(model_path)
+        model = EncoderDecoderModel.from_pretrained(model_path).cuda()
+        predictor = Predictor(model, tokenizer, device="cuda", num_titles=5)
+        print(f'Ready - running on GPU.')
+    else:
+        print(f'No GPU available - using hosted inference API.')
+        predictor = HostedInference(model_path, num_titles=5)
     
-    predictor = Predictor(model, tokenizer, device=device, num_titles=5)
     interface = create_gradio_ui(predictor)
     interface.launch()