stal76 commited on
Commit
d1e1d59
·
1 Parent(s): 4000803

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -2
app.py CHANGED
@@ -19,10 +19,15 @@ class Net(nn.Module):
19
  def forward(self,x):
20
  return self.layer(x)
21
 
22
- @st.cache(allow_output_mutation=True)
23
- def GetModelAndTokenizer():
24
  model = Net()
25
  model.load_state_dict(torch.load('model.dat', map_location=torch.device('cpu')))
 
 
 
 
 
26
  tokenizer = AutoTokenizer.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
27
  model_emb = AutoModelForSeq2SeqLM.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
28
  return model, tokenizer, model_emb
 
19
  def forward(self,x):
20
  return self.layer(x)
21
 
22
+ @st.cache
23
+ def GetModel():
24
  model = Net()
25
  model.load_state_dict(torch.load('model.dat', map_location=torch.device('cpu')))
26
+ return model
27
+
28
+ #@st.cache(allow_output_mutation=True)
29
+ def GetModelAndTokenizer():
30
+ model = GetModel()
31
  tokenizer = AutoTokenizer.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
32
  model_emb = AutoModelForSeq2SeqLM.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
33
  return model, tokenizer, model_emb