from transformers import AutoTokenizer, AutoModelWithLMHead, SummarizationPipeline
import gradio as gr
import tokenize
import io

pipeline = SummarizationPipeline(
    model=AutoModelWithLMHead.from_pretrained("SEBIS/code_trans_t5_large_source_code_summarization_python_multitask_finetune"),
    tokenizer=AutoTokenizer.from_pretrained("SEBIS/code_trans_t5_large_source_code_summarization_python_multitask_finetune", skip_special_tokens=True),
    device=0)
    
def code_summarizer(code):
  def code_tokenizer(line):
    result= []
    line = io.StringIO(line) 
    
    for toktype, tok, start, end, line in tokenize.generate_tokens(line.readline):
      if (not toktype == tokenize.COMMENT):
        if toktype == tokenize.STRING:
          result.append("CODE_STRING")
        elif toktype == tokenize.NUMBER:
          result.append("CODE_INTEGER")
        elif (not tok=="\n") and (not tok=="    "):
          result.append(str(tok))
    return ' '.join(result)
    
  tokenized_code = code_tokenizer(code)
  summary = pipeline(tokenized_code)
  return summary[0]['summary_text']

def call_examples():
  examples = ['''def findAverage(list): sum = 0. for x in list: sum = sum + x average = sum / len(list) return average''',
              '''def findMax(list): max = list[0] for x in list: if x > max: max = x return max''',
              '''def findMin(list): min = list[0] for x in list: if x < min: min = x return min''']
  return examples

gr.Interface(fn=code_summarizer,
             inputs=gr.inputs.Textbox(
                 lines=5,
                 default='',
                 placeholder='Insert a Python code here',
                 label='PYTHON CODE'),
             outputs=gr.outputs.Textbox(
                 type='auto',
                 label='CODE SUMMARY'),
             title='Code Summarizer From CodeTrans',
             description='Summarize any Python code',
             examples=call_examples(),
             allow_flagging='never').launch(inbrowser=True)