Separate GPU Process

app.py

@@ -77,8 +77,16 @@ def get_prompt(schema, query_input):
     """
     return text
 
+@spaces.GPU(duration=60)
+def generate_sql(prompt):
+
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    input_token_len = input_ids.shape[1]
+    outputs = model.generate(input_ids.to(model.device), max_new_tokens=1024)
+    result = tokenizer.decode(outputs[0][input_token_len:], skip_special_tokens=True)
+    return result
+
 # Generate SQL
-@spaces.GPU(duration=120)
 def text2sql(table, query_input):
     if table is None:
         return {
@@ -94,10 +102,7 @@ def text2sql(table, query_input):
     print(f'Prompt Generated...')
     try:
         print(f'Generating SQL... {model.device}')
-        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-        input_token_len = input_ids.shape[1]
-        outputs = model.generate(input_ids.to(model.device), max_new_tokens=1024)
-        result = tokenizer.decode(outputs[0][input_token_len:], skip_special_tokens=True)
+        result = generate_sql(prompt)
         print('SQL Generated...')
     except Exception as e:
         return {