tools/summarize_news.py

from smolagents import Tool
from typing import Any, Optional

class SimpleTool(Tool):
    name = "summarize_news"
    description = "T\u00f3m t\u1eaft n\u1ed9i dung v\u0103n b\u1ea3n ti\u1ebfng Vi\u1ec7t t\u1eeb b\u00e0i b\u00e1o."
    inputs = {'text': {'type': 'string', 'description': 'Văn bản tiếng Việt cần tóm tắt.'}}
    output_type = "string"

    def forward(self, text: str) -> str:
        """
        Tóm tắt nội dung văn bản tiếng Việt từ bài báo.

        Args:
            text (str): Văn bản tiếng Việt cần tóm tắt.

        Returns:
            str: Văn bản đã được tóm tắt.
        """
        import torch
        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model_name = "VietAI/vit5-base-vietnews-summarization"

        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
        model = model.to(device)

        formatted_text = "vietnews: " + text + " </s>"
        encoding = tokenizer(formatted_text, return_tensors="pt")
        input_ids = encoding["input_ids"].to(device)
        attention_masks = encoding["attention_mask"].to(device)

        with torch.no_grad():
            outputs = model.generate(
                input_ids=input_ids,
                attention_mask=attention_masks,
                max_length=256,
            )

        summary = tokenizer.decode(
            outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True
        )
        return summary