relik-ie commited on
Commit
11ecbac
·
verified ·
1 Parent(s): ee57afb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -6
app.py CHANGED
@@ -175,8 +175,12 @@ relik_models = {
175
  "sapienzanlp/relik-entity-linking-large": Relik.from_pretrained(
176
  "sapienzanlp/relik-entity-linking-large",
177
  device="cuda",
178
- index=wikipedia_index,
179
- retriever=wikipedia_retriever,
 
 
 
 
180
  reader_kwargs={"dataset_kwargs": {"use_nme": True}},
181
  ),
182
  "relik-ie/relik-cie-small": Relik.from_pretrained(
@@ -207,15 +211,23 @@ relik_models = {
207
  ),
208
  "relik-ie/relik-relation-extraction-small": Relik.from_pretrained(
209
  "relik-ie/relik-relation-extraction-small",
210
- index=wikidata_index,
 
 
211
  device="cuda",
212
- retriever=wikidata_retriever,
 
 
213
  ),
214
  "relik-ie/relik-entity-linking-large-robust": Relik.from_pretrained(
215
  "relik-ie/relik-entity-linking-large-robust",
216
- index=wikipedia_index,
 
 
217
  device="cuda",
218
- retriever=wikipedia_retriever,
 
 
219
  reader_kwargs={"dataset_kwargs": {"use_nme": True}},
220
  ),
221
  }
 
175
  "sapienzanlp/relik-entity-linking-large": Relik.from_pretrained(
176
  "sapienzanlp/relik-entity-linking-large",
177
  device="cuda",
178
+ index={
179
+ TaskType.SPAN: wikipedia_index,
180
+ },
181
+ retriever={
182
+ TaskType.SPAN:wikipedia_retriever,
183
+ },
184
  reader_kwargs={"dataset_kwargs": {"use_nme": True}},
185
  ),
186
  "relik-ie/relik-cie-small": Relik.from_pretrained(
 
211
  ),
212
  "relik-ie/relik-relation-extraction-small": Relik.from_pretrained(
213
  "relik-ie/relik-relation-extraction-small",
214
+ index={
215
+ TaskType.TRIPLET:wikidata_index,
216
+ },
217
  device="cuda",
218
+ retriever={
219
+ TaskType.TRIPLET: wikidata_retriever,
220
+ },
221
  ),
222
  "relik-ie/relik-entity-linking-large-robust": Relik.from_pretrained(
223
  "relik-ie/relik-entity-linking-large-robust",
224
+ index={
225
+ TaskType.SPAN: wikipedia_index,
226
+ },
227
  device="cuda",
228
+ retriever={
229
+ TaskType.SPAN: wikipedia_retriever,
230
+ },
231
  reader_kwargs={"dataset_kwargs": {"use_nme": True}},
232
  ),
233
  }