prateekagrawal commited on
Commit
db188a6
1 Parent(s): 8477933

Fixed error in ai4bharat

Browse files
Files changed (1) hide show
  1. apps/inference.py +4 -4
apps/inference.py CHANGED
@@ -14,12 +14,12 @@ def load_model(masked_text, model_name):
14
  if model_name == "flax-community/roberta-hindi":
15
  from_flax = True
16
 
17
- if model_name in [
18
- "ai4bharat/indic-bert",
19
- "neuralspace-reverie/indic-transformers-hi-bert",
20
- ]:
21
  masked_text = masked_text.replace("<mask>", "[MASK]")
22
 
 
 
 
23
  st.write(model_name, masked_text)
24
 
25
  model = AutoModelForMaskedLM.from_pretrained(model_name, from_flax=from_flax)
 
14
  if model_name == "flax-community/roberta-hindi":
15
  from_flax = True
16
 
17
+ if model_name == "neuralspace-reverie/indic-transformers-hi-bert":
 
 
 
18
  masked_text = masked_text.replace("<mask>", "[MASK]")
19
 
20
+ elif model_name == "ai4bharat/indic-bert":
21
+ masked_text = masked_text.replace("<mask>", "([MASK])")
22
+
23
  st.write(model_name, masked_text)
24
 
25
  model = AutoModelForMaskedLM.from_pretrained(model_name, from_flax=from_flax)