codeexplainer / app.py
umarfarooqtoru's picture
Update app.py
2c583db verified
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
def load_model():
model_name = "Salesforce/codet5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
return pipeline("summarization", model=model, tokenizer=tokenizer)
def explain_code(files):
explainer = load_model()
explanations = []
for file in files:
filename = file.name if hasattr(file, 'name') else 'Uploaded File'
if isinstance(file, bytes):
code = file.decode('utf-8')
elif hasattr(file, 'read'):
code = file.read().decode('utf-8')
else:
code = str(file)
lines = code.splitlines()
line_explanations = []
for idx, line in enumerate(lines):
if not line.strip():
continue
input_length = len(line.split())
max_length = max(1, input_length // 2)
result = explainer(line, max_length=max_length, do_sample=False)[0].get('summary_text', '')
line_explanations.append(f"Line {idx+1}: {line}\nExplanation: {result.strip()}\n")
explanations.append(f"## {filename}\n" + "\n".join(line_explanations))
return "\n\n".join(explanations)
def main():
with gr.Blocks() as demo:
gr.Markdown("# CodeExplainerBot (Local CodeT5-small)")
code_input = gr.File(label="Upload code files", file_count="multiple", type="binary")
output = gr.Textbox(label="Line-by-line Explanation", lines=20)
explain_btn = gr.Button("Explain Code")
explain_btn.click(fn=explain_code, inputs=[code_input], outputs=output)
demo.launch()
if __name__ == "__main__":
main()