paultltc commited on
Commit
017336c
·
1 Parent(s): d3703b7

refactor -- following guidelines to push tools

Browse files
Files changed (1) hide show
  1. tool.py +3 -0
tool.py CHANGED
@@ -1,3 +1,5 @@
 
 
1
  from dataclasses import dataclass
2
  from typing import List, Optional, Tuple
3
 
@@ -154,6 +156,7 @@ class VisualRAGTool(Tool):
154
  output_type = "string"
155
 
156
  model_name: str = "vidore/colqwen2-v1.0"
 
157
 
158
  def _init_models(self, model_name: str) -> None:
159
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
1
+ import os
2
+
3
  from dataclasses import dataclass
4
  from typing import List, Optional, Tuple
5
 
 
156
  output_type = "string"
157
 
158
  model_name: str = "vidore/colqwen2-v1.0"
159
+ api_key: str = os.getenv("OPENAI_KEY")
160
 
161
  def _init_models(self, model_name: str) -> None:
162
  self.device = "cuda" if torch.cuda.is_available() else "cpu"