| | import streamlit as st |
| | from transformers import AutoModelForSeq2SeqLM, AutoTokenizer |
| | import torch |
| |
|
| | @st.cache_resource |
| | def load_model(): |
| | tokenizer = AutoTokenizer.from_pretrained("google/mt5-base", padding_side="left", use_fast=False) |
| | model = AutoModelForSeq2SeqLM.from_pretrained("google/mt5-base") |
| | return tokenizer, model |
| |
|
| | st.title("Український Чат-бот") |
| |
|
| | if "history" not in st.session_state: |
| | st.session_state.history = [] |
| |
|
| | if "user_input" not in st.session_state: |
| | st.session_state.user_input = "" |
| |
|
| | tokenizer, model = load_model() |
| |
|
| | def send_message(): |
| | if st.session_state.user_input: |
| | inputs = tokenizer(st.session_state.history + [st.session_state.user_input], return_tensors="pt", padding=True, truncation=True) |
| | with torch.no_grad(): |
| | outputs = model.generate(**inputs, max_length=100) |
| | response = tokenizer.decode(outputs[0], skip_special_tokens=True) |
| | st.session_state.history.extend([st.session_state.user_input, response]) |
| | st.session_state.user_input = "" |
| |
|
| | def update_user_input(): |
| | st.session_state.user_input = st.session_state.temp_user_input |
| |
|
| | st.text_input("Ви:", key="temp_user_input", on_change=update_user_input) |
| |
|
| | if st.button("Надіслати"): |
| | send_message() |
| |
|
| | |
| | if st.session_state.get("temp_user_input") and st.session_state.get("last_input", "") != st.session_state.get("temp_user_input"): |
| | st.session_state["last_input"] = st.session_state["temp_user_input"] |
| | send_message() |
| |
|
| | if st.session_state.history: |
| | for i in range(0, len(st.session_state.history), 2): |
| | st.write(f"Ви: {st.session_state.history[i]}") |
| | if i + 1 < len(st.session_state.history): |
| | st.write(f"Бот: {st.session_state.history[i+1]}") |
