import gradio as gr
from transformers import pipeline
from gradio.themes.utils import sizes


theme = gr.themes.Default(radius_size=sizes.radius_none).set(
    block_label_text_color = '#4D63FF',
    block_title_text_color = '#4D63FF',
    button_primary_text_color = '#4D63FF',
    button_primary_background_fill='#FFFFFF',
    button_primary_border_color='#4D63FF',
    button_primary_background_fill_hover='#EDEFFF',
)


sentimentPipeline = pipeline('zero-shot-classification', model='roberta-large-mnli')

def sentiment_analysis(text, labels):
    candidate_labels = labels.split(',') 
    results = sentimentPipeline(text, candidate_labels)
    total_results = ""
    index = 0

    for candidate_label in candidate_labels:
        total_results += f"Sentiment: {results.get('labels')[index]}, Score: {results.get('scores')[index]}"
        total_results += '\r\n'
        index += 1

    return total_results
    
demo = gr.Interface(fn=sentiment_analysis, 
                    inputs=[
                        gr.components.Textbox(label="Text"),
                        gr.components.Textbox(label="Label")
                    ],
                    outputs='text',
                    examples=[['I am happy', 'negative, netural, positive'], ['I am sad', 'negative, netural, positive']],
                    theme = theme,
                    title = "文本情感分析"
                    )


if __name__ == "__main__":
    demo.queue(concurrency_count=10)
    demo.launch(server_name = "0.0.0.0")