Update app.py
Browse files
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
|
23 |
-
def
|
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
|