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)