from transformers import pipeline
import streamlit as st

# Definition der Chatbots
chatbots = {
    "GPT-Neo 2.7B": "EleutherAI/gpt-neo-2.7B",
    "DialoGPT Small": "microsoft/DialoGPT-small"
}

# Funktion zum Abrufen des ausgewählten Chatbots
def get_chatbot(name):
    return pipeline('text-generation', model=chatbots[name])

# Definiere die Streamlit-Oberfläche
st.title("Prompt Generator")
chatbot_name = st.selectbox("Wähle einen Chatbot aus:", list(chatbots.keys()))
prompt = st.text_input("Gib einen Satzanfang ein:")
button = st.button("Generieren")

# Generiere neuen Text beim Klick auf den Button
if button:
    chatbot = get_chatbot(chatbot_name)
    text = chatbot(prompt, max_length=30, do_sample=True, temperature=0.7)[0]['generated_text']
    st.write(text)