Spaces:
Sleeping
Sleeping
File size: 1,718 Bytes
439fcd4 c7a8772 439fcd4 c7a8772 439fcd4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import torch
from transformers import AutoTokenizer, EncoderDecoderModel
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_name = "Shuu12121/CodeEncoderDecodeerModel-Ghost"
# Tokenizerの読み込み
encoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/encoder_tokenizer")
decoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/decoder_tokenizer")
if decoder_tokenizer.pad_token is None:
decoder_tokenizer.pad_token = decoder_tokenizer.eos_token
model = EncoderDecoderModel.from_pretrained(model_name).to(device)
model.eval()
def generate_docstring(code: str) -> str:
inputs = encoder_tokenizer(
code,
return_tensors="pt",
padding=True,
truncation=True,
max_length=2048
).to(device)
with torch.no_grad():
output_ids = model.generate(
input_ids=inputs.input_ids,
attention_mask=inputs.attention_mask,
max_length=256,
num_beams=5,
early_stopping=True,
decoder_start_token_id=model.config.decoder_start_token_id,
eos_token_id=model.config.eos_token_id,
pad_token_id=model.config.pad_token_id,
no_repeat_ngram_size=2
)
return decoder_tokenizer.decode(output_ids[0], skip_special_tokens=True)
# Gradio UI
import gradio as gr
iface = gr.Interface(
fn=generate_docstring,
inputs=gr.Textbox(label="Code Snippet", lines=10, placeholder="Paste your function here..."),
outputs=gr.Textbox(label="Generated Docstring"),
title="Code-to-Docstring Generator",
description="This demo uses a custom encoder-decoder model to generate docstrings from code."
)
iface.launch()
|