Models/HFAgent.py

from transformers import pipeline

# MODEL_CKPT = "HuggingFaceH4/zephyr-7b-beta"
MODEL_CKPT = "AVeryRealHuman/DialoGPT-small-TonyStark"

class HFAgent:
    def __init__(self):
        self.pipe = pipeline("conversational", MODEL_CKPT)
    
    def generate(self, chat_history):
        return self.pipe(chat_history)
    
    def __call__(self, chat_history):
        return self.generate(chat_history)
    
    def __repr__(self):
        return f"HFAgent(model={self.pipe.model})"
    
    def __str__(self):
        return f"HFAgent(model={self.pipe.model})"
    

## For testing purposes
# def main():
#     agent = HFAgent()
#     messages = [
#         {
#             "role": "system",
#             "content": "You are a friendly chatbot who always responds in the style of a pirate",
#         },
#         {"role": "user", "content": "How many hotdogs can a human eat in one sitting?"},
#     ]
#     new_messages = agent(messages)
#     print(new_messages[-1])

# if __name__ == "__main__":
#     main()