Pretrain-GPT / app.py
ahmadmac's picture
Update app.py
f7b8f0c verified
raw
history blame
1.75 kB
import streamlit as st
import os
from transformers import AutoTokenizer, AutoModelForCausalLM
st.set_page_config(page_title="GPT-2 Text Generator", layout="centered")
#tokenizer = AutoTokenizer.from_pretrained("ahmadmac/Pretrained-GPT2")
#model = AutoModelForCausalLM.from_pretrained("ahmadmac/Pretrained-GPT2")
tokenizer = AutoTokenizer.from_pretrained("ahmadmac/DistillGPT2-CSV")
model = AutoModelForCausalLM.from_pretrained("ahmadmac/DistillGPT2-CSV")
#google_api_key=
import google.generativeai as genai
GOOGLE_API_KEY=os.environ["google_api_key"]
genai.configure(api_key=GOOGLE_API_KEY)
# def generate_text(prompt):
# inputs = tokenizer(prompt, return_tensors="pt")
# outputs = model.generate(**inputs, max_length=50)
# generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
# return generated_text
gemini_model = genai.GenerativeModel('gemini-1.5-pro')
def generate_text(input_text):
input_ids = tokenizer.encode(input_text, return_tensors="pt")
trained_output = model.generate(input_ids, max_length=100, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
trained_response = tokenizer.decode(trained_output[0], skip_special_tokens=True)
prompt = f"Improve this text to make it clearer and more concise: {trained_response}"
generated_text = gemini_model.generate_content(prompt)
return generated_text
st.title("GPT-2 Text Generator")
st.write("Enter a prompt to generate text using GPT-2")
user_input = st.text_input("Prompt")
if st.button("Generate"):
if user_input:
with st.spinner("Generating..."):
generated_text = generate_text(user_input)
st.write(generated_text.text)
else:
st.warning("Please enter a prompt")