ManishThota commited on
Commit
a5beb04
·
verified ·
1 Parent(s): 98b0c4e

Update src/text_processor.py

Browse files
Files changed (1) hide show
  1. src/text_processor.py +3 -1
src/text_processor.py CHANGED
@@ -4,6 +4,7 @@ from pydantic import BaseModel
4
  import json
5
  import warnings
6
  import spaces
 
7
 
8
  # Ignore warnings
9
  warnings.filterwarnings(action='ignore')
@@ -34,6 +35,7 @@ generation_args = {
34
  "do_sample": True
35
  }
36
 
 
37
  class LLMHelper:
38
  def __init__(self, pipeline):
39
  self.chatbot = pipeline
@@ -80,7 +82,7 @@ class VideoAnalysis(BaseModel):
80
  # Create an instance of LLMHelper (using the already loaded pipeline)
81
  llm_helper = LLMHelper(pipe)
82
 
83
- def process_llm_output(description -> Dict:
84
  # Generate the logic from the LLM output
85
  generated_logic = llm_helper.generate_logic(description
86
 
 
4
  import json
5
  import warnings
6
  import spaces
7
+ from typing import Dict
8
 
9
  # Ignore warnings
10
  warnings.filterwarnings(action='ignore')
 
35
  "do_sample": True
36
  }
37
 
38
+ @spaces.GPU(duration=75)
39
  class LLMHelper:
40
  def __init__(self, pipeline):
41
  self.chatbot = pipeline
 
82
  # Create an instance of LLMHelper (using the already loaded pipeline)
83
  llm_helper = LLMHelper(pipe)
84
 
85
+ def process_llm_output(description) -> Dict:
86
  # Generate the logic from the LLM output
87
  generated_logic = llm_helper.generate_logic(description
88