File size: 386 Bytes
1f6a5b9 |
1 2 3 4 5 6 7 8 9 10 11 |
# front_end_agent.py
class FrontEndAgent:
def __init__(self, model, tokenizer):
self.model = model
self.tokenizer = tokenizer
def process(self, task_data):
inputs = self.tokenizer(task_data['task'], return_tensors='pt')
outputs = self.model.generate(**inputs)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|