import streamlit as st
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM

st.title("Jogi Dokumentumok AI Tokenizálása és Válaszgenerálás")

# Információs szekció
st.write(
    """
    Ez az alkalmazás tokenizálja a megadott jogi dokumentumot az általad választott modell segítségével,
    majd generál válaszokat a megadott modell használatával.
    """
)

# Tokenizáló modell választása
tokenizer_name = st.selectbox("Tokenizáló Modell", ["pile-of-law/legalbert-large-1.7M-2", "bert-base-uncased", "bert-base-multilingual-cased"])
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

# Válaszgeneráló modell választása
generator_name = st.selectbox("Válaszgeneráló Modell", ["microsoft/phi-2", "gpt2", "gpt2-medium"])
model = AutoModelForCausalLM.from_pretrained(generator_name, trust_remote_code=True)

# Input mező a jogi dokumentumhoz
document_text = st.text_area("Jogi Dokumentum", "")

# Generálási paraméterek testreszabása
max_length = st.slider("Maximális Válasz Hossz", min_value=10, max_value=200, value=50, step=10)
num_return_sequences = st.slider("Válaszok Száma", min_value=1, max_value=5, value=3, step=1)
temperature = st.slider("Hőmérséklet", min_value=0.1, max_value=2.0, value=0.7, step=0.1)

if st.button("Tokenizálás"):
    # Tokenizálás
    tokenized_document = tokenizer(document_text, return_tensors="pt")

    # Tokenizált dokumentum mentése
    with open("tokenized_document.txt", "w", encoding="utf-8") as file:
        file.write(str(tokenized_document))

if st.button("Válaszgenerálás"):
    # Válasz generálása
    input_ids = tokenized_document["input_ids"]
    output = model.generate(input_ids, max_length=max_length, num_return_sequences=num_return_sequences, temperature=temperature)

    # Válaszok megjelenítése és mentése fájlba
    st.subheader("Generált Válaszok:")
    
    # Válaszok kiválasztása
    selected_responses = st.multiselect("Válaszok Kiválasztása", [f"Válasz {i+1}" for i in range(num_return_sequences)])
    
    for i, generated_sequence in enumerate(output):
        response_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
        
        # Válasz megjelenítése csak akkor, ha kiválasztották
        if f"Válasz {i+1}" in selected_responses:
            st.write(f"Válasz {i+1}: {response_text}")
            
            # Valószínűségek megjelenítése
            if hasattr(model, "get_output_embeddings"):
                logits = model.get_output_embeddings()(generated_sequence)
                probs = logits.softmax(dim=-1)
                st.write("Valószínűségek:")
                st.table(probs.tolist())

            # Fájlba mentés opció
            if st.button(f"Mentés Válasz {i+1}"):
                with open(f"generated_response_{i+1}.txt", "w", encoding="utf-8") as file:
                    file.write(response_text)
                st.write(f"A válasz {i+1} sikeresen el lett mentve.")

    # Generált válaszok egyetlen fájlba mentése
    if st.button("Mentés Összes Válasz"):
        with open("all_generated_responses.txt", "w", encoding="utf-8") as file:
            for i, generated_sequence in enumerate(output):
                generated_text = tokenizer.decode(generated_sequence, skip_special_tokens=True)
                file.write(f"Válasz {i+1}:\n")
                file.write(generated_text + "\n\n")
        st.write("Az összes válasz sikeresen el lett mentve egy fájlba.")