import torch
import torch.nn as nn
from transformers import AutoModel
from huggingface_hub import PyTorchModelHubMixin
from transformers import AutoTokenizer

# Define the model:
BASE_MODEL = "Peltarion/xlm-roberta-longformer-base-4096"
class ReadabilityModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self, model_name=BASE_MODEL):
        super(ReadabilityModel, self).__init__()
        self.model = AutoModel.from_pretrained(model_name)
        self.drop = nn.Dropout(p=0.2)
        self.fc = nn.Linear(768, 1)

    def forward(self, ids, mask):
        out = self.model(input_ids=ids, attention_mask=mask,
                         output_hidden_states=False)
        out = self.drop(out[1])
        outputs = self.fc(out)

        return outputs
    
# Load the model:
model = ReadabilityModel.from_pretrained("trokhymovych/TRank_readability")

# Load the tokenizer:
tokenizer = AutoTokenizer.from_pretrained("trokhymovych/TRank_readability")

# Set the model to evaluation mode
model.eval()
# Example input text
input_text = "This is an example sentence to evaluate readability."
# Tokenize the input text
inputs = tokenizer.encode_plus(
    input_text,
    add_special_tokens=True,
    max_length=512,
    truncation=True,
    padding='max_length',
    return_tensors='pt'
)
ids = inputs['input_ids']
mask = inputs['attention_mask']

# Make prediction
with torch.no_grad():
    outputs = model(ids, mask)
    readability_score = outputs.item()
    
# Print the input text and the readability score
print(f"Input Text: {input_text}")
print(f"Readability Score: {readability_score}")

@misc{trokhymovych2024openmultilingualscoringreadability,
      title={An Open Multilingual System for Scoring Readability of Wikipedia}, 
      author={Mykola Trokhymovych and Indira Sen and Martin Gerlach},
      year={2024},
      eprint={2406.01835},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2406.01835}, 
}