added decorator to generate action (#69)

* added decorator to generate action * changed decoration to prepare_model
sotopia-lab · Apr 29, 2024 · d762753 · d762753
1 parent cad4701
commit d762753
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sotopia_generate.py b/sotopia_generate.py
@@ -3,7 +3,6 @@
 from typing import TypeVar
 from functools import cache
 import logging
-import json
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
@@ -38,6 +37,7 @@
 log = logging.getLogger("generate")
 logging_handler = LoggingCallbackHandler("langchain")
 
+# @spaces.GPU
 def generate_action(
     model_name: str,
     history: str,
@@ -82,6 +82,7 @@ def generate_action(
     #     print(e)
     #     return AgentAction(action_type="none", argument="")
 
+@spaces.GPU(duration=1200)
 @cache
 def prepare_model(model_name):
     compute_type = torch.float16