import gradio as gr
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForSequenceClassification


sample_text = [
    [
        "Poems on various subjects. Whereto is prefixed a short essay on the structure of English verse"
    ],
    [
        "Journal of a Residence in China and the neighbouring countries from 1830 to 1833. With an introductory essay by the Hon. and Rev. Baptist Wriothesley Noel. [With a map.]"
    ],
    ["The Adventures of Oliver Twist. [With plates.]"],
    ["['The Adventures of Sherlock Holmes', 'Single Works']"],
    [
        "['Coal, Iron, and Oil; or, the Practical American miner. A plain and popular work on our mines and mineral resources ... With numerous maps and engravings, etc']"
    ],
    [
        "Summer Travelling in Iceland; being the narrative of two journeys across the island ... With a chapter on Askja by E. Delmar Morgan ... Containing also a literal translation of three sagas. Maps, etc'"
    ],
    [
        "Histoire de France au moyen aÃÇge, depuis Philippe-Auguste jusqu'aÃÄ la fin du reÃÄgne de Louis XI. 1223-1483. Troisieme eÃÅdition"
    ],
    [
        "Two Centuries of Soho: its institutions, firms, and amusements. By the Clergy of St. Anne's, Soho, J. H. Cardwell ... H. B. Freeman ... G. C. Wilton ... assisted by other contributors, etc"
    ],
    ["""A Christmas Carol"""],
]

description = """
British Library Books genre detection model
"""

tokenizer = AutoTokenizer.from_pretrained("BritishLibraryLabs/bl-books-genre")

model = AutoModelForSequenceClassification.from_pretrained("BritishLibraryLabs/bl-books-genre")
classifier = pipeline('text-classification',model=model, tokenizer=tokenizer, return_all_scores=True)


def predict(text):
    predictions = classifier(text)
    return {pred['label']: pred['score'] for pred in predictions[0]}

gr.Interface(predict,
    inputs=gr.inputs.Textbox(label="Book title"),
    outputs=gr.outputs.Label(label="Predicted genre"), 
    interpretation='shap',
    examples=sample_text,description=description,
   ).launch(enable_queue=True)