File size: 1,517 Bytes
bcf0ac0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import gradio as gr
from union.remote import UnionRemote
from flytekit.types.file import FlyteFile
remote = UnionRemote()
# Fetch recent executions
recent_executions = remote.recent_executions(limit=10)
executions = [
e for e in recent_executions if e.spec.launch_plan.name == "bert-fine-tune.bert_ft"
]
recent_ex_id = executions[0].id.name
execution = remote.fetch_execution(name=recent_ex_id)
ft_llm = execution.outputs["o0"].remote_source
model_cache = execution.outputs["o1"].remote_source
predict_task = remote.fetch_task(name="bert-fine-tune.predict_sentiment")
def execute_flyte_task(text):
inputs = {
"text": text,
"model_cache_dir": model_cache,
"model": FlyteFile(ft_llm)
}
execution = remote.execute(predict_task, inputs=inputs, wait=True)
# Extract the response
response = execution.outputs["o0"]
# Format the response
sentiment = response['label']
score = response['score']
# Determine color based on sentiment
color = "red" if sentiment == "NEGATIVE" else "green"
# Format the output HTML
output_html = f"""
<div style="text-align: center;">
<h2>Sentiment: <span style="color: {color};">{sentiment}</span></h2>
<p>Confidence Score: {score:.2f}</p>
</div>
"""
return output_html
# Launch Gradio app
iface = gr.Interface(
fn=execute_flyte_task,
inputs=["text"],
outputs=gr.HTML(), # Change output to HTML for better formatting
live=False,
)
iface.launch(debug=True)
|