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")