egmaminta's picture
Update app.py
77a5389
raw
history blame
2.26 kB
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=-1)
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 findRange(list):
return max(list)-min(list)''']]
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',
theme='dark-peach',
examples=call_examples(),
allow_flagging='never').launch(inbrowser=True)