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

# Load your fine-tuned CamemBERT NER model
model_name = "GEODE/camembert-base-edda-span-classification"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)


ner_pipeline = pipeline("token-classification", model=model, tokenizer=tokenizer, aggregation_strategy="simple")

def extract_coordinates(text):
    # Run NER
    entities = ner_pipeline(text)
    
    # Collect coordinate entities
    coords_text = [ent['word'] for ent in entities if ent['entity_group'] == 'Latlong']
    
    if coords_text:
        return "\n".join(coords_text)
    else:
        return "No coordinates found"


def norm_coordinates(text):

    generator = pipeline("text2text-generation", model="GEODE/mt5-small-coords-norm")
    
    # Example input text
    input_text = "extract_coordinates: " + text

    # Generate prediction using the pipeline
    predicted_coordinates_from_pipeline = generator(input_text, max_length=128)

    return predicted_coordinates_from_pipeline[0]['generated_text']

examples = [
    "* AACH ou ACH, s. f. petite ville d'Allemagne dans le cercle de Souabe, près de la source de l'Aach. Long. 26. 57. lat. 47. 55.",
    "* ARCALU (Principauté d') petit état des Tartares-Monguls, sur la riviere d'Hoamko, où commence la grande muraille de la Chine, sous le 122e degré de longitude & le 42e de latitude septentrionale.",
    "* ARÉQUIPE, ou ARIQUIPA, (Géog.) ville de l'Amérique méridion. dans le Pérou, sur une riviere, dans un terrein fertile. Long. 308. lat. mérid. 16. 40.",
    "* AUTAN-KELURAN, (Géog.) ville du Turquestan. Long. 110d. & lat. 46. 45. selon Uluhbeg ; & long. 116. & lat. 45. selon Nassiredden.",
    "Boston ; c'est le nom qu'on a donné à la ville capitale de la nouvelle Angleterre, dans l'Amérique septentrionale ; elle est grande & a un très-bon port. Lat. 42 degrés, 20 minutes ; long. 306 degrés, 50 & quelques minutes.",
    "CABEÇA-DE-VIDE, (Géog.) petite ville avec château, en Portugal, dans l'Alentéjo, à cinq lieues de Port-Alegre. Longitude 10. 48. latitude 39."
]


with gr.Blocks() as demo:
    gr.Markdown("## Geographic Coordinate Extractor")
    
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(
                label="Enter text", 
                placeholder="e.g. * AACH ou ACH, s. f. petite ville d'Allemagne dans le cercle de Souabe, près de la source de l'Aach. Long. 26. 57. lat. 47. 55.",
                lines=3
            )

            run_btn = gr.Button("Extract & Show")
            
    
        with gr.Column():
            out_text = gr.Textbox(label="Extracted coordinates (fine-tuned CamemBERT NER)")
            run_btn.click(fn=extract_coordinates, inputs=inp, outputs=out_text)

            out_text = gr.Textbox(label="Extracted and normalized coordinates (fine-tuned mT5)")
            run_btn.click(fn=norm_coordinates, inputs=inp, outputs=out_text)

    with gr.Row():
        
        gr.Examples(
            examples=examples,
            inputs=inp,
            label="Examples"
        )

# Launch
if __name__ == "__main__":
    demo.launch()