languageModel / app.py
nastasiasnk's picture
Update app.py
f51b7a0 verified
raw
history blame contribute delete
No virus
1.27 kB
import streamlit as st
import os
from transformers import pipeline, set_seed
from huggingface_hub import HfFolder
import transformers
import torch
# 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 identifie
#generator = pipeline("text-generation", model="openai-community/gpt2")
generator = pipeline('text-generation', model='gpt2-large')
st.title("Text Generation")
st.write("Enter your text below.")
text = st.text_area("Your input")
st.write("Enter seed.")
seed_input = st.text_area("Set seed")
st.write("Enter max length.")
maxLength = st.text_area("max length")
# Convert seed input to integer
try:
seed = int(seed_input)
max_length = int(maxLength)
except ValueError:
seed = 1
max_length = 100
set_seed(seed)
if st.button("Generate Text"):
# Use default values or handle None appropriately
if seed is not None:
set_seed(seed)
if text and max_length:
# Generate text
out = generator(text, max_length=max_length, num_return_sequences=1)
st.json(out)
st.write(f"Reply: {out[0]['generated_text']}")