File size: 1,582 Bytes
4567b01 10e9b7d 31243f4 4567b01 3c4371f 4567b01 e80aab9 3c4371f 4567b01 |
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 |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
class BasicAgent:
def __init__(self):
print("Loading model...")
model_id = "google/flan-t5-small"
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model.to(self.device)
print("BasicAgent with FLAN-T5 initialized.")
def __call__(self, question: str) -> str:
prompt = (
"You are a general AI assistant. I will ask you a question. "
"Report your thoughts, and finish your answer with the following template: "
"FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible "
"OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number "
"neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, "
"neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. "
f"Now, here is the question: {question}"
)
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True).to(self.device)
outputs = self.model.generate(**inputs, max_new_tokens=100)
answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"Answer: {answer}")
return answer |