chatbot.py

from transformers import pipeline

chat_pipeline = pipeline("text-generation", model="distilgpt2")

def get_response(prompt):
    response = chat_pipeline(prompt, max_new_tokens=50, do_sample=True)
    return response[0]['generated_text']