import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import nltk

# Wczytywanie modelu i tokenizera
model_name = "clui/opus-mt-it-pl"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Upewnij się, że zasoby NLTK są zainstalowane
nltk.download('punkt')

def preprocess_text(input_text):
    """Usuwa nagłówki i zbędne spacje z tekstu."""
    return ' '.join(line.strip() for line in input_text.splitlines() if line.strip() and not line.startswith('['))

def translate_text(text):
    """Tłumaczy tekst przy użyciu modelu."""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).input_ids
    outputs = model.generate(inputs, max_new_tokens=200, do_sample=True, top_k=30, top_p=0.95)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def translate_long_text(long_text):
    """Tłumaczy długi tekst, dzieląc go na całe zdania."""
    sentences = nltk.sent_tokenize(long_text)  # Dzielimy tekst na zdania
    translated_sentences = [translate_text(sentence) for sentence in sentences]  # Tłumaczymy każde zdanie
    return ' '.join(translated_sentences)  # Łączymy przetłumaczone zdania

# Interfejs użytkownika
st.title("Tłumacz")
text = st.text_area("Wprowadź tekst:")

if st.button("Tłumacz"):
    if text:
        clean_text = preprocess_text(text)
        if clean_text:
            translated_text = translate_long_text(clean_text)
            st.subheader("Przetłumaczone:")
            st.write(translated_text)
        else:
            st.warning("Brak tekstu do przetłumaczenia po przetworzeniu.")
    else:
        st.warning("Proszę wprowadzić tekst do przetłumaczenia.")