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

# Define the Hugging Face repository paths for each model
tokenizer_repo = "viv/UD_Greek-GUD"
lemmatizer_repo = "viv/UD_Greek-GUD"  # Adjust this path if necessary
tagger_repo = "viv/UD_Greek-GUD"  # Adjust this path if necessary
parser_repo = "viv/UD_Greek-GUD"  # Adjust this path if necessary

# Load models using the Hugging Face model hub
tokenizer = AutoTokenizer.from_pretrained(tokenizer_repo)
lemmatizer_model = torch.load("models/el_new_nocharlm_lemmatizer.pt")
tagger_model = torch.load("models/el_new_transformer_tagger.pt")
parser_model = torch.load("models/el_new_transformer_parser.pt")

# Prediction function
def predict(text):
    # Tokenize input
    inputs = tokenizer(text, return_tensors="pt")
    
    # Perform lemmatization
    lemma_outputs = lemmatizer_model(**inputs)
    lemmas = lemma_outputs.logits.argmax(-1).tolist()  # Process lemmatizer output
    
    # Perform POS tagging
    pos_outputs = tagger_model(**inputs)
    pos_tags = pos_outputs.logits.argmax(-1).tolist()  # Process tagger output
    
    # Perform dependency parsing
    dep_outputs = parser_model(**inputs)
    dep_parse = dep_outputs.logits.argmax(-1).tolist()  # Process parser output
    
    # Return results
    return {
        "lemmas": lemmas,
        "pos_tags": pos_tags,
        "dep_parse": dep_parse,
    }

# Gradio Interface
interface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="json",
    title="Greek NLP Pipeline",
    description="Perform lemmatization, POS tagging, and dependency parsing for Greek text using custom models.",
)

# Launch interface
interface.launch()