Update app.py

app.py

@@ -7,12 +7,28 @@ st.set_page_config(page_title="GPT-2 Text Generator", layout="centered")
 #model = AutoModelForCausalLM.from_pretrained("ahmadmac/Pretrained-GPT2")
 tokenizer = AutoTokenizer.from_pretrained("ahmadmac/Pretrain-CSV-GPT2")
 model = AutoModelForCausalLM.from_pretrained("ahmadmac/results")
+#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 
+
 
-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
 
 st.title("GPT-2 Text Generator")
 st.write("Enter a prompt to generate text using GPT-2")
@@ -23,6 +39,6 @@ if st.button("Generate"):
     if user_input:
         with st.spinner("Generating..."):
             generated_text = generate_text(user_input)
-            st.write(generated_text)
+            st.write(generated_text.text)
     else:
         st.warning("Please enter a prompt")