Alikestocode commited on
Commit
a217627
·
1 Parent(s): 9a4d6d3

Fix: Remove context manager usage for spaces.GPU decorator

Browse files

- spaces.GPU is a decorator, not a context manager
- Use decorator with maximum duration (1800s) to allow user flexibility
- Keep gpu_duration parameter for user awareness/UI
- Fix AttributeError: __enter__ error

Files changed (1) hide show
  1. app.py +13 -10
app.py CHANGED
@@ -291,9 +291,7 @@ def _generate_router_plan_streaming_internal(
291
  yield "", {}, f"❌ Invalid model choice: {model_choice}. Available: {list(MODELS.keys())}", ""
292
  return
293
 
294
- # Use GPU context manager with user-specified duration
295
- with spaces.GPU(duration=gpu_duration):
296
- try:
297
  prompt = build_router_prompt(
298
  user_task=user_task,
299
  context=context,
@@ -382,14 +380,14 @@ def _generate_router_plan_streaming_internal(
382
  parsed_plan = {}
383
  validation_msg = f"❌ JSON parsing failed: {exc}"
384
 
385
- yield completion, parsed_plan, validation_msg, prompt
386
-
387
- except Exception as exc:
388
- error_msg = f"❌ Generation failed: {str(exc)}"
389
- yield "", {}, error_msg, ""
390
 
391
 
392
- @spaces.GPU(duration=600) # Default wrapper for backward compatibility
393
  def generate_router_plan_streaming(
394
  user_task: str,
395
  context: str,
@@ -403,7 +401,12 @@ def generate_router_plan_streaming(
403
  top_p: float,
404
  gpu_duration: int = 600,
405
  ):
406
- """Wrapper function that calls internal generator with GPU duration."""
 
 
 
 
 
407
  yield from _generate_router_plan_streaming_internal(
408
  user_task, context, acceptance, extra_guidance,
409
  difficulty, tags, model_choice, max_new_tokens,
 
291
  yield "", {}, f"❌ Invalid model choice: {model_choice}. Available: {list(MODELS.keys())}", ""
292
  return
293
 
294
+ try:
 
 
295
  prompt = build_router_prompt(
296
  user_task=user_task,
297
  context=context,
 
380
  parsed_plan = {}
381
  validation_msg = f"❌ JSON parsing failed: {exc}"
382
 
383
+ yield completion, parsed_plan, validation_msg, prompt
384
+
385
+ except Exception as exc:
386
+ error_msg = f"❌ Generation failed: {str(exc)}"
387
+ yield "", {}, error_msg, ""
388
 
389
 
390
+ @spaces.GPU(duration=1800) # Use maximum duration to allow user flexibility
391
  def generate_router_plan_streaming(
392
  user_task: str,
393
  context: str,
 
401
  top_p: float,
402
  gpu_duration: int = 600,
403
  ):
404
+ """
405
+ Generate router plan with streaming output.
406
+
407
+ Note: gpu_duration parameter is for user awareness. The actual GPU allocation
408
+ uses the decorator's duration (1800s max) to allow flexibility.
409
+ """
410
  yield from _generate_router_plan_streaming_internal(
411
  user_task, context, acceptance, extra_guidance,
412
  difficulty, tags, model_choice, max_new_tokens,