Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModel, pipeline | |
| import torch | |
| MODEL_NAME = "microsoft/codebert-base" | |
| # ๅๅงๅ tokenizer ๅ model | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModel.from_pretrained(MODEL_NAME) | |
| # ๅปบ็ซไธๅ feature-extraction pipeline | |
| feature_extractor = pipeline( | |
| task="feature-extraction", | |
| model=model, | |
| tokenizer=tokenizer, | |
| return_tensors=True, # ่ฎ่ผธๅบไฟๆ tensor ๅฝขๅผ | |
| framework="pt" # ๆๅฎ PyTorch | |
| ) | |
| def embed_text(input_text): | |
| # pipeline ้ ่จญ่ผธๅบ็ๆฏ List[List[List[float]]] ็ไธ็ถญ้ฃๅ | |
| # ๅฏไพ้ๆฑๅ pooling | |
| outputs = feature_extractor(input_text) | |
| # ๅ่จญๅชๅ็ฌฌไธๅ batch ไธฆๅๅนณๅๆฑ ๅ | |
| # outputs[0] => List[List[float]] | |
| # ๅ ถไธญๆฏๅ token ้ฝๆฏไธๅๅ้ | |
| token_embeddings = torch.tensor(outputs[0]) | |
| embedding = torch.mean(token_embeddings, dim=0) # ๅๆๆ token ็ๅนณๅ | |
| return embedding.tolist() | |
| # ็จ Gradio ๅปบ็ซไธๅไป้ข | |
| with gr.Blocks(): | |
| gr.Markdown("# Feature Extraction Demo with CodeBERT") | |
| text_input = gr.Textbox(label="Enter text here") | |
| output = gr.JSON(label="Vector Embedding") | |
| btn = gr.Button("Compute Embedding") | |
| btn.click(embed_text, inputs=text_input, outputs=output) | |
| gr.Examples( | |
| examples=["def hello_world(): print('Hello World!')", "print('test')"], | |
| inputs=text_input, | |
| outputs=output | |
| ) | |
| gr.launch() | |