import streamlit as st
from huggingface_hub import InferenceClient
from langchain_community.document_loaders import ArxivLoader

st.title("📝 Arxiv Summaries mit Zephyr 7B Beta")

# Sidebar optional, falls du hier ggf. einen Hugging Face Token abfragen möchtest
with st.sidebar:
    st.markdown("## Einstellungen")
    hf_api_token = st.text_input("Hugging Face API Token (optional)", type="password")

# Eingabe für den Arxiv-Suchbegriff
arxiv_query = st.text_input(
    "Gib einen Suchbegriff für Arxiv ein",
    placeholder="Beispiel: 'reasoning'",
)

# Button zum Ausführen
if st.button("Arxiv durchsuchen und zusammenfassen"):
    if not arxiv_query:
        st.warning("Bitte einen Suchbegriff eingeben.")
    else:
        try:
            # 1) ArxivLoader initialisieren
            loader = ArxivLoader(
                query=arxiv_query,
                load_max_docs=2
            )
            
            # 2) Dokumente laden (Summaries)
            docs = loader.get_summaries_as_docs()
            if not docs:
                st.warning("Keine Dokumente gefunden.")
            else:
                # Wir nehmen hier exemplarisch nur das erste Dokument
                doc_content = docs[0].page_content

                # 3) Verbindung zum Zephyr-7B-Beta-Modell
                # Falls du einen Token verwendest, kannst du den hier übergeben:
                # client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=hf_api_token)
                client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

                # Prompt erstellen: Kontext = Inhalt der Summaries
                prompt = f"""
Bitte fasse den folgenden Inhalt prägnant zusammen (auf Deutsch):

{doc_content}

Gib mir eine gut strukturierte Zusammenfassung der Studien.
"""

                # 4) Anfrage an das Modell (2048 Tokens als Richtwert)
                response = client.text_generation(
                    prompt,
                    max_new_tokens=2048,
                    # Weitere Parameter nach Bedarf, z.B.:
                    # temperature=0.7,
                    # do_sample=True,
                )

                # Ausgabe in Streamlit
                st.write("### Zusammenfassung")
                st.write(response)

        except Exception as e:
            st.error(f"Fehler beim Laden oder Zusammenfassen: {e}")