test / app.py
fxcb's picture
Update app.py
e2fc846
raw
history blame contribute delete
523 Bytes
import streamlit as st
from transformers import AutoModelForMaskedLM
from transformers import AutoTokenizer, FillMaskPipeline
from pprint import pprint
tokenizer_hf = AutoTokenizer.from_pretrained('projecte-aina/roberta-base-ca-v2')
model = AutoModelForMaskedLM.from_pretrained('projecte-aina/roberta-base-ca-v2')
model.eval()
pipeline = FillMaskPipeline(model, tokenizer_hf)
text = f"Tinc pressa, marxo al MACBA a veure l'exposició <mask>."
res_hf = pipeline(text)
res = ([r['token_str'] for r in res_hf])
st.write(res)