import gradio as gr
import torch
import pandas as pd
from transformers import AutoModelForSequenceClassification, BertTokenizerFast

#Setting up the model
model_save_path = "Kelmoir/2026-02-14_CAFA6_classification_Rostlab_prot_bert_v1"
model = AutoModelForSequenceClassification.from_pretrained(model_save_path)
tokenizer = BertTokenizerFast.from_pretrained("Rostlab/prot_bert",
                                              do_lower_case=False)

# Setting up the computation device - cuda or CPU
device = "cuda" if torch.cuda.is_available() else "cpu"

# Get the id2label
id2label = model.config.id2label

model.to(device)

def prepare_for_tokenization(input_prot_seq: str)-> str:
  """
  This function shall insert a [SPACE] between each protein sequence item, and returns the resulting string,
  also, uppercases the sting and replaces unlike proteins
  A maximum length of 1024 is enforced as was used for model training
  """
  if len(input_prot_seq) > 1024:
    input_prot_seq = input_prot_seq[:1024]
  input_prot_seq = input_prot_seq.strip().upper().replace(r"[UZOB]", "X")
  spaces = " "*len(input_prot_seq)

# Source - https://stackoverflow.com/a
# Posted by Ma0, modified by community. See post 'Timeline' for change history
# Retrieved 2026-01-12, License - CC BY-SA 3.0
  return ''.join(map(''.join, zip(input_prot_seq, spaces))).strip()  # kudos @Coldspeed

def process_result(logits, threshold=0.6, top_n=10):
  """
  This function will take in a single protein logit prediction and turns it into alist of predicted GO-terms
  Args:
  **logits** a single Logit array
  **threshold** A float value between 0 and 1. Only GO-terms with a score higher than this value will be considered.
  **top_n** An integer value. If no GO-terms are above the threshold, the top_n GO-terms will be returned.
  Ensure, that the length is equal
  """
  all_results = []
  # Convert logits to probabilities using sigmoid
  probabilities = torch.sigmoid(logits).squeeze()

  # Get indices of GO terms above threshold
  above_threshold_indices = torch.where(probabilities > threshold)[0]

  # If no terms above threshold, take the top_n terms
  if len(above_threshold_indices) == 0:
    top_n_values, top_n_indices = torch.topk(probabilities, k=min(top_n, len(probabilities)))
    selected_indices = top_n_indices
    selected_probabilities = top_n_values
  else:
    selected_indices = above_threshold_indices
    selected_probabilities = probabilities[above_threshold_indices]

  # Sort selected terms by probability in descending order
  sorted_probabilities, sort_indices = torch.sort(selected_probabilities, descending=True)
  sorted_indices = selected_indices[sort_indices]

  for idx, prob in zip(sorted_indices, sorted_probabilities):
    go_term = id2label[idx.item()]
    score = prob.item()
    all_results.append([go_term, f"{score:.3f}"]) # Format score to 3 decimal places

  results_df = pd.DataFrame(all_results)
  # The request specifies 3 columns, all string, no headers, no index
  return results_df

def predict_on_input(input: str, threshold:float):
  model.eval()
  with torch.inference_mode():
    tokenized = tokenizer(prepare_for_tokenization(input),
                          return_tensors="pt").to(device)
    output_logits = model(**tokenized)

  output_table = process_result(output_logits.logits,
                                threshold=threshold)

  return output_table

description = """
This is a demo project to showcase my attempts at predicting the protein functions - known as GO-terms based of the protein amino acid string.

This is essentially the scope of the CAFA 6 challenge that was hosted on Kaggle - https://www.kaggle.com/competitions/cafa-6-protein-function-prediction/overview

This demo takes in a single protein sequence, and will then predict the corresonding GO-terms based of that, and display the score.

Right now, no further post-processing or frills are available.

The model used to perform the predictions is a fine tuned ProtBERT Model. See: https://huggingface.co/Rostlab/prot_bert

And finally, the model isn't to good right now, there is much more experimentation ahead. For instance, the example was reported with the following GO-terms:
- GO:1990837
- GO:0005515
"""

demo = gr.Interface(
    fn = predict_on_input,
    inputs = [
        gr.Textbox(value="MHHRMNEMNLSPVGMEQLTSSSVSNALPVSGSHLGLAASPTHSAIPAPGLPVAIPNLGPSLSSLPSALSLMLPMGIGDRGVMCGLPERNYTLPPPPYPHLESSYFRTILPGILSYLADRPPPQYIHPNSINVDGNTALSITNNPSALDPYQSNGNVGLEPGIVSIDSRSVNTHGAQSLHPSDGHEVALDTAITMENVSRVTSPISTDGMAEELTMDGVAGEHSQIPNGSRSHEPLSVDSVSNNLAADAVGHGGVIPMHGNGLELPVVMETDHIASRVNGMSDSALSDSIHTVAMSTNSVSVALSTSHNLASLESVSLHEVGLSLEPVAVSSITQEVAMGTGHVDVSSDSLSFVSPSLQMEDSNSNKENMATLFTIWCTLCDRAYPSDCPEHGPVTFVPDTPIESRARLSLPKQLVLRQSIVGAEVGVWTGETIPVRTCFGPLIGQQSHSMEVAEWTDKAVNHIWKIYHNGVLEFCIITTDENECNWMMFVRKARNREEQNLVAYPHDGKIFFCTSQDIPPENELLFYYSRDYAQQIGVPEHPDVHLCNCGKECNSYTEFKAHLTSHIHNHLPTQGHSGSHGPSHSKERKWKCSMCPQAFISPSKLHVHFMGHMGMKPHKCDFCSKAFSDPSNLRTHLKIHTGQKNYRCTLCDKSFTQKAHLESHMVIHTGEKNLKCDYCDKLFMRRQDLKQHVLIHTQERQIKCPKCDKLFLRTNHLKKHLNSHEGKRDYVCEKCTKAYLTKYHLTRHLKTCKGPTSSSSAPEEEEEDDSEEEDLADSVGTEDCRINSAVYSADESLSAHK",
                   label="Protein sequence",
                   show_label=True),
        gr.Slider(minimum = -1,
                  maximum = 1,
                  step= 0.01,
                  value=0.25,
                  label="Prediction threshold",
                  show_label=True)
    ],
    outputs = [
        gr.Dataframe(label="Outpt table of Go-terms",
                     show_label=True)
    ],
    description = description,
    title ="CAFA ProtBERT prediction demo V1",
)
#demo.launch(debug=True)
demo.launch()