Update app.py

app.py

@@ -1,40 +1,23 @@
-
-
 import streamlit as st
 import os
-import torch  # Add this line to import the torch library
-import transformers
-
-
-HF_TOKEN = os.getenv('HF_TOKEN')
-from huggingface_hub import HfFolder
-# Set the token using HfFolder (this persists the token)
-HfFolder.save_token(HF_TOKEN)
-
 from transformers import pipeline
+from huggingface_hub import HfFolder
 
-# Load the model, specifying the use of GPU if available
-
-#device = 0 if torch.cuda.is_available() else -1  # use GPU if available
-
-#generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B', device=device)
-
-
-generator = pipeline("text-generation", model="gpt-2")
+# Ensure the HF_TOKEN environment variable is set correctly
+HF_TOKEN = os.getenv('HF_TOKEN')
+if HF_TOKEN:
+    HfFolder.save_token(HF_TOKEN)
+else:
+    st.warning("HF_TOKEN is not set. Proceeding without a token.")
 
+# Use a valid model identifier
+generator = pipeline("text-generation", model="gpt2")
 
-st.title("text class")
-st.write("your text.")
-text = st.text_area("your input")
+st.title("Text Generation")
+st.write("Enter your text below.")
+text = st.text_area("Your input")
 
 if text:
     out = generator(text, do_sample=False)
     st.json(out)
-    st.write(f"reply: {out}")
-
-
-
-
-
-
-
+    st.write(f"Reply: {out[0]['generated_text']}")