import gradio as gr
import os, vertexai

from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

from vertexai.language_models import TextGenerationModel

#credentials = os.environ["CREDENTIALS"]
#project = os.environ["PROJECT"]

#
from utils import authenticate
credentials, project = authenticate()
#

vertexai.init(project = project, 
              location = "us-west1", 
              credentials = credentials)

generation_model = TextGenerationModel.from_pretrained("text-bison@001")

def invoke(prompt):
    return generation_model.predict(prompt = prompt).text

gr.close_all()
demo = gr.Interface(fn=invoke, 
                    inputs = [gr.Textbox(label = "Prompt", lines = 1)],
                    outputs = [gr.Textbox(label = "Completion", lines = 1)],
                    title = "Generative AI - Text",
                    description = "<a href='https://www.gradio.app/'>Gradio</a> UI using <a href='https://cloud.google.com/vertex-ai?hl=en/'>Google Vertex AI</a> API with Bison foundation model")
demo.launch()