import os

os.system("pip install awscli boto3 sagemaker gradio")

os.system("aws configure set aws_access_key_id AKIAYIBNUHKT7UN3PKCK")
os.system("aws configure set aws_secret_access_key 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")
os.system("aws configure set default.region us-east-1")
os.system("aws configure set default.output json")


import boto3
import sagemaker
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri
import gradio as gr

predictor = None

def deploy_and_get_interface(hf_model_url):
    global predictor



    try:
        role = sagemaker.get_execution_role()
    except ValueError:
        iam = boto3.client('iam')
        role = iam.get_role(RoleName='Moh-work')['Role']['Arn']

    hf_model_id = hf_model_url.split('/')[-2]

    hub = {
        'HF_MODEL_ID': hf_model_id,
        'SM_NUM_GPUS': '1'
    }

    huggingface_model = HuggingFaceModel(
        image_uri=get_huggingface_llm_image_uri("huggingface", version="1.0.3"),
        env=hub,
        role=role
    )

    predictor = huggingface_model.deploy(
        initial_instance_count=1,
        instance_type="ml.g5.2xlarge",
        container_startup_health_check_timeout=300,
    )

    def get_prediction(input_text):
        response = predictor.predict({"inputs": input_text})
        return response[0]['generated_text']

    iface = gr.Interface(fn=get_prediction, 
                         inputs="text", 
                         outputs="text",
                        
                         title="HuggingFace Model Predictor",
                         description="Enter some text and get the model's prediction!")

    return iface

def get_prediction(input_text):
    response = predictor.predict({"inputs": input_text})
    return response[0]['generated_text']

def trigger_deploy(hf_model_url):
    deploy_and_get_interface(hf_model_url)
    iface = gr.Interface(fn=get_prediction, 
                         inputs="text", 
                         outputs="text",
                         
                         title="HuggingFace Model Predictor",
                         description="Enter some text and get the model's prediction!")
    iface.launch()

def setup_interface():
    iface = gr.Interface(fn=trigger_deploy,
                     inputs="text", 
                     outputs="text",
                     title="Enter HuggingFace Model URL",
                     description="Please input the URL of the desired HuggingFace model and submit to deploy.")
    iface.launch()

setup_interface()