SandeepU/code-explainer-c

1from transformers import AutoTokenizer, T5ForConditionalGeneration2import torch3 4def load_model():5    model_name = "Salesforce/codet5-base-multi-sum"6    tokenizer = AutoTokenizer.from_pretrained(model_name)7    model = T5ForConditionalGeneration.from_pretrained(model_name)8    model.eval()9    model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))10    return tokenizer, model11 12def generate_explanation(code, tokenizer, model):13    device = model.device14    # Final prompt style: generate docstring15    input_text = f"generate docstring: {code.strip()}"16    input_ids = tokenizer.encode(input_text, return_tensors="pt", truncation=True).to(device)17    output = model.generate(input_ids, max_new_tokens=150, early_stopping=True)18    return tokenizer.decode(output[0], skip_special_tokens=True)19