如何在 python 中获取序列的上下文？

Question

Given a text such as

The green ball on the table is mine

I want to extract the contexts of size 2 (to the right and left) of any sequence in the sentence above.

For example, given the sequence green ball , i expect to have "The", "on","the,"table" as the contexts.

Also given The , I want to have "green", "ball" "on", "the" as the contexts. Given given mine , I want to have "on", "the" "table", "is" as the contexts.

Since The does not have contexts on the left. it picks the 4 contexts from the right, the same applies to mine .

I tried out something but my method was based on splitting as shown below

query = "green ball"
context_window=2
texts="The green ball  on the table is mine"

tokens = texts.split()
index = tokens.index(query)
begin = max(0, index - context_window)
end = min(index + 1 + context_window, len(tokens))
context_words = tokens[begin:end]

I discovered it cannot work in this case. Any way out?

Answer 1

Split the input string into 2 parts: the part to the left of the query and the part to the right of it.

Then split the left and right parts into words. Get the last context_window words of the left part and the first context_window words of the right part. If either part is shorter than context_window , add the difference to the number of words you get from the other part.

Finally, get the last words from the left part and the first words from the right part.

query = "green ball"
context_window=2
texts="The green ball  on the table is mine"
left = texts[0:texts.find(query)]
right = texts[texts.find(query)+len(query):]

if left:
    left_tokens = left.split()
else:
    left_tokens = []

if right:
    right_tokens = right.split()
else:
    right_tokens = []

if len(left_tokens) >= context_window and len(right_tokens) >= context_window:
    left_context = context_window
    right_context = context_window
elif len(left_tokens) < context_window:
    left_context = len(left_tokens)
    right_context = context_window + (context_window - left_context)
else:
    right_context = len(right_tokens)
    left_context = context_window + (context_window - right_context)

context_words = left_tokens[-left_context:] + right_tokens[0:right_context]