#using pipeline to predict the input text
from transformers import pipeline, AutoTokenizer
import torch

label_mapping = {
    'delete': [0, 'LABEL_0'],
    'keep': [1, 'LABEL_1'],
    'merge': [2, 'LABEL_2'],
    'no consensus': [3, 'LABEL_3'],
    'speedy keep': [4, 'LABEL_4'],
    'speedy delete': [5, 'LABEL_5'],
    'redirect': [6, 'LABEL_6'],
    'withdrawn': [7, 'LABEL_7']
}

def predict_text(text, model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = pipeline("text-classification", model=model_name, return_all_scores=True)
    
    # Tokenize and truncate the text
    tokens = tokenizer(text, truncation=True, max_length=512)
    truncated_text = tokenizer.decode(tokens['input_ids'], skip_special_tokens=True)
    
    results = model(truncated_text)
    final_scores = {key: 0.0 for key in label_mapping}
    
    for result in results[0]:
        for key, value in label_mapping.items():
            if result['label'] == value[1]:
                final_scores[key] = result['score']
                break
    
    return final_scores


# from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
# import torch

# label_mapping = {
#     'delete': [0, 'LABEL_0'],
#     'keep': [1, 'LABEL_1'],
#     'merge': [2, 'LABEL_2'],
#     'no consensus': [3, 'LABEL_3'],
#     'speedy keep': [4, 'LABEL_4'],
#     'speedy delete': [5, 'LABEL_5'],
#     'redirect': [6, 'LABEL_6'],
#     'withdrawn': [7, 'LABEL_7']
# }

# def predict_text(text, model_name):
#     tokenizer = AutoTokenizer.from_pretrained(model_name)
#     model = AutoModelForSequenceClassification.from_pretrained(model_name, output_attentions=True)
    
#     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
#     outputs = model(**inputs)
#     predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    
#     final_scores = {key: 0.0 for key in label_mapping}
#     for i, score in enumerate(predictions[0]):
#         for key, value in label_mapping.items():
#             if i == value[0]:
#                 final_scores[key] = score.item()
#                 break
    
#     # Calculate average attention
#     attentions = outputs.attentions
#     avg_attentions = torch.mean(torch.stack(attentions), dim=1)  # Average over all layers
#     avg_attentions = avg_attentions.mean(dim=1)[0]  # Average over heads
#     token_importance = avg_attentions.mean(dim=0)
    
#     # Decode tokens and highlight important ones
#     tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
#     highlighted_text = []
#     for token, importance in zip(tokens, token_importance):
#         if importance > token_importance.mean():
#             highlighted_text.append(f"<b>{token}</b>") #
#         else:
#             highlighted_text.append(token)
    
#     highlighted_text = " ".join(highlighted_text)
#     highlighted_text = highlighted_text.replace("##", "") 
    
#     return final_scores, highlighted_text