Update app.py

app.py

@@ -1,34 +1,25 @@
 import streamlit as st
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
-import torch
+from transformers import pipeline
 
-# Initialisation de Streamlit
-st.title("Générateur de texte avec Hugging Face")
-st.write("Entrez un prompt pour générer du texte.")
+# Titre de l'application
+st.title("Générateur de Texte avec Hugging Face")
 
-# Chargement du modèle et du tokenizer
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-model = GPT2LMHeadModel.from_pretrained("gpt2")
+# Sélection du modèle
+model_name = st.selectbox("Choisissez un modèle :", ["gpt2", "distilgpt2"])
 
-def generate_text(prompt):
-    # Tokenisation de l'entrée
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
-
-    # Génération de texte
-    outputs = model.generate(inputs, max_length=50, num_return_sequences=1)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return generated_text
+# Chargement du pipeline de génération de texte
+text_generator = pipeline("text-generation", model=model_name)
 
 # Entrée utilisateur
-user_input = st.text_input("Votre prompt:", "")
+prompt = st.text_area("Entrez votre prompt :", "Écrivez ici...")
 
+# Bouton pour générer le texte
 if st.button("Générer"):
-    if user_input:
+    if prompt:
         with st.spinner("Génération en cours..."):
-            generated = generate_text(user_input)
-            st.success("Texte généré:")
-            st.write(generated)
+            generated_text = text_generator(prompt, max_length=50, num_return_sequences=1)[0]['generated_text']
+            st.success("Texte généré :")
+            st.write(generated_text)
     else:
         st.error("Veuillez entrer un prompt.")