hotel_bot / app.py
Daddario's picture
Update app.py
4266281 verified
raw
history blame contribute delete
898 Bytes
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}")