| from smolagents import Tool | |
| from typing import Any, Optional | |
| class SimpleTool(Tool): | |
| name = "summarize_news" | |
| description = "Summarizes the given Vietnamese news text." | |
| inputs = {"text":{"type":"string","description":"The Vietnamese news text to be summarized."}} | |
| output_type = "string" | |
| def forward(self, text: str) -> str: | |
| """ | |
| Summarizes the given Vietnamese news text. | |
| Args: | |
| text (str): The Vietnamese news text to be summarized. | |
| Returns: | |
| str: The summarized version of the input text. | |
| """ | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import torch | |
| model_name = "VietAI/vit5-base-vietnews-summarization" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.bfloat16) | |
| model.cuda() | |
| formatted_text = "vietnews: " + text + " </s>" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| 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 |