Simple Keyword Extraction using OpenAIChatGenerator

!pip install haystack-ai

import os
from getpass import getpass

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = getpass("Enter OpenAI API key:")

import json


from haystack.dataclasses import ChatMessage
from haystack.components.builders import ChatPromptBuilder
from haystack.components.generators.chat import OpenAIChatGenerator

text_to_analyze = "Artificial intelligence models like large language models are increasingly integrated into various sectors including healthcare, finance, education, and customer service. They can process natural language, generate text, translate languages, and extract meaningful insights from unstructured data. When performing key word extraction, these systems identify the most significant terms, phrases, or concepts that represent the core meaning of a document. Effective extraction must balance between technical terminology, domain-specific jargon, named entities, action verbs, and contextual relevance. The process typically involves tokenization, stopword removal, part-of-speech tagging, frequency analysis, and semantic relationship mapping to prioritize terms that most accurately capture the document's essential information and main topics."

messages = [
    ChatMessage.from_user(
        '''
You are a keyword extractor. Extract the most relevant keywords and phrases from the following text. For each keyword:
1. Find single and multi-word keywords that capture important concepts
2. Include the starting position (index) where each keyword appears in the text
3. Assign a relevance score between 0 and 1 for each keyword
4. Focus on nouns, noun phrases, and important terms

Text to analyze: {{text}}

Return the results as a JSON array in this exact format:
{
  "keywords": [
    {
      "keyword": "example term",
      "positions": [5],
      "score": 0.95
    },
    {
      "keyword": "another keyword",
      "positions": [20],
      "score": 0.85
    }
  ]
}

Important:
- Each keyword must have its EXACT character position in the text (counting from 0)
- Scores should reflect the relevance (0–1)
- Include both single words and meaningful phrases
- List results from highest to lowest score
'''
    )
]

builder = ChatPromptBuilder(template=messages, required_variables='*')
prompt = builder.run(text=text_to_analyze)

# Initialize the chat-based generator
extractor = OpenAIChatGenerator(model="gpt-4o-mini")

# Run the generator with our formatted prompt
results = extractor.run(
    messages=prompt["prompt"],
    generation_kwargs={"response_format": {"type": "json_object"}}
)

# Extract the raw text reply
output_str = results["replies"][0].text

try:
    data = json.loads(output_str)
    for kw in data["keywords"]:
        print(f'Keyword: {kw["keyword"]}')
        print(f' Positions: {kw["positions"]}')
        print(f' Score: {kw["score"]}\n')
except json.JSONDecodeError:
    print("Failed to parse the output as JSON. Raw output:", output_str)