import streamlit as st
from transformers import BertTokenizer, BertForTokenClassification
from transformers import pipeline

# Carica il modello e il tokenizer
model = BertForTokenClassification.from_pretrained("./hotel_model")
tokenizer = BertTokenizer.from_pretrained("./hotel_model")

# Funzione di inferenza
def predict_entities(text):
    # Tokenizza il testo
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    # Ottieni le predizioni dal modello
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_ids = torch.argmax(logits, dim=-1)
    # Converti le predizioni in etichette (potresti aver bisogno di una mappatura delle etichette)
    return predicted_ids

st.title("Hotel Bot")
query = st.text_input("Inserisci una query:")

if query:
    entities = predict_entities(query)
    st.write(f"Entità estratte: {entities}")