Data enrichment

enrichment.py

llm = ChatOpenAI(max_tokens=200)
newsgroups_train_enriched = []

data = newsgroups_train.data
for index in range(len(data)):
    # ... (Data preprocessing)
    messages = [
        SystemMessage(content="You are a helpful assistant that classifies texts."),
        HumanMessage(content=f"Give one category to the following text: {text}. The category should be among this list: {categories}")
    ]
    response = llm(messages)
    newsgroups_train_enriched.append(text + response.content)