import streamlit as st
from setfit import SetFitModel

# Load the model
model = SetFitModel.from_pretrained("peter2000/vulnerable-groups-setfit")

# Define the classes
group_dict = {
    1: 'Women and girls',
 2: 'Children and youth',
 3: 'Landlocked countries',
 4: 'Outdoor workers',
 5: 'Riverine and flood-prone areas',
 6: 'Small-scale farmers',
 7: 'Men and boys',
 8: 'Small island developing states (SIDS)',
 9: 'Fisherfolk and fishing communities',
 10: 'Children with disabilities',
 11: 'Low-income households',
 12: 'Rural communities',
 13: 'Pregnant women and new mothers',
 14: 'Young adults',
 15: 'Urban slums',
 16: 'Gender non-conforming individuals',
 17: 'Remote communities',
 18: 'Older adults and the elderly',
 19: 'Elderly population',
 20: 'Mountain communities',
 21: 'People with disabilities',
 22: 'Indigenous peoples',
 23: 'Informal settlements and slums',
 24: 'Coastal communities',
 25: 'Informal sector workers',
 26: 'Drought-prone regions',
 27: 'People with pre-existing health conditions',
 28: 'Small-scale farmers and subsistence agriculture',
 29: 'Migrants and displaced populations',
 30: 'no vulnerable group mentioned'}

def predict(text):
    preds = model([text])[0].item()
    return group_dict[preds]

text = st.text_area('enter your text here')

x = st.slider('Select a value')
st.write(x, 'squared is', x * x)