Guy24 commited on
Commit
ed75380
·
1 Parent(s): 3902836

adding application

Browse files
Files changed (1) hide show
  1. app.py +10 -0
app.py CHANGED
@@ -4,6 +4,15 @@ from huggingface_hub import login
4
  # run once at startup
5
  if "HF_TOKEN" in os.environ:
6
  login(token=os.environ["HF_TOKEN"])
 
 
 
 
 
 
 
 
 
7
  import torch
8
  from exceptiongroup import catch
9
  from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -117,6 +126,7 @@ class PatchscopesRetriever(WordRetrieverBase):
117
 
118
  return last_token_hidden_states
119
 
 
120
  def get_hidden_states_and_retrieve_word(self, word, num_tokens_to_generate=None):
121
  last_token_hidden_states = self.extract_hidden_states(word)
122
  patchscopes_description_by_layers = self.retrieve_word(
 
4
  # run once at startup
5
  if "HF_TOKEN" in os.environ:
6
  login(token=os.environ["HF_TOKEN"])
7
+
8
+ # app.py
9
+ import os; os.environ.setdefault('HF_HOME', '/data/hf-cache')
10
+ os.environ.setdefault('HF_HUB_ENABLE_HF_TRANSFER', '1')
11
+
12
+ from huggingface_hub import login
13
+ login(os.getenv("HF_TOKEN", ""))
14
+
15
+ from spaces import GPU
16
  import torch
17
  from exceptiongroup import catch
18
  from transformers import AutoTokenizer, AutoModelForCausalLM
 
126
 
127
  return last_token_hidden_states
128
 
129
+ @GPU # this block runs on a job GPU
130
  def get_hidden_states_and_retrieve_word(self, word, num_tokens_to_generate=None):
131
  last_token_hidden_states = self.extract_hidden_states(word)
132
  patchscopes_description_by_layers = self.retrieve_word(