import gradio as gr
import google.generativeai as palm
import os

# Get your own API key on https://developers.generativeai.google/products/palm
api_key = os.environ["api_key"]
palm.configure(api_key=api_key)

defaults = {
    'model': 'models/text-bison-001',
    'temperature': 0.7,
    'candidate_count': 1,
    'top_k': 40,
    'top_p': 0.95,
    'max_output_tokens': 1024,
    'stop_sequences': [],
    'safety_settings': [
        {"category": "HARM_CATEGORY_DEROGATORY", "threshold": 3},
        {"category": "HARM_CATEGORY_TOXICITY", "threshold": 3},
        {"category": "HARM_CATEGORY_VIOLENCE", "threshold": 3},
        {"category": "HARM_CATEGORY_SEXUAL", "threshold": 3},
        {"category": "HARM_CATEGORY_MEDICAL", "threshold": 3},
        {"category": "HARM_CATEGORY_DANGEROUS", "threshold": 3}]

}

with open("example.txt") as f:
    example = f.read()

with gr.Blocks() as app:
    def chat(text):
        response = palm.generate_text(
        **defaults,
        prompt = f"""Rewrite the following sentence to fix the grammar issues and correct the sentence.
                {example}
                input: {text}
                fixed"""
        )
        return response.result

    with gr.Column():
        text = gr.Textbox(lines=4, label="Text", max_lines=4, placeholder="Write something awesome. It will be corrected automatically.")

    with gr.Column():
        output = gr.Textbox(lines=4, label="Output", max_lines=4, show_copy_button=True)
        
        with gr.Row():
            clr_btn = gr.ClearButton([text, output], variant="primary")
            btn = gr.Button("Submit")
            btn.click(fn=chat, inputs=text, outputs=output)


if __name__ == "__main__":
    app.launch()