# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

using MultiResolutionIterators, CorpusLoaders
include("../src/ner.jl")

using LinearAlgebra, Statistics

dataset = flatten_levels(collect(CorpusLoaders.load(GMB())) , lvls(GMB, :document)) |> full_consolidate

X = [word.(sentence) for sentence in dataset]
Y = [CorpusLoaders.named_entity.(sentence) for sentence in dataset]

unique(vcat(Y...))

ner = NERTagger()

# +
function try_outs(ner, x_in, y_in)
    unique_labels = unique(ner.model.labels)
    num_labels = length(unique_labels)
    confusion_matrix = zeros(Int, (num_labels, num_labels))

    for (x_seq, y_seq) in zip(x_in, y_in)
        preds = ner(x_seq)

        for (pred, logit) in zip(preds, y_seq)
            pred == "MISC" && continue

            if(logit == "O")
                confusion_matrix[findfirst(x -> x==pred, unique_labels), findfirst(x -> x=="O", unique_labels)] += 1
            elseif(logit == "Location")
                confusion_matrix[findfirst(x -> x==pred, unique_labels), findfirst(x -> x=="LOC", unique_labels)] += 1
            elseif(logit == "Person")
                confusion_matrix[findfirst(x -> x==pred, unique_labels), findfirst(x -> x=="PER", unique_labels)] += 1
            elseif(logit == "Organization")
                confusion_matrix[findfirst(x -> x==pred, unique_labels), findfirst(x -> x=="ORG", unique_labels)] += 1
            else
                continue
            end
        end
    end

#     print(confusion_matrix)
    s1 = sum(confusion_matrix, dims=2)
    s2 = sum(confusion_matrix, dims=1)
    dg = diag(confusion_matrix)
    s1 = [s1[1:2]..., s1[4:5]...]
    s2 = [s2[1:2]..., s2[4:5]...]
    dg = [dg[1:2]..., dg[4:5]...]

    a = mean(dg ./ s1)
    b = mean(dg ./ s2)

    println("Precision and recall are:", a, " ", b)
    println("F1 is:", (2 * a * b) / (a + b))
end

# -

try_outs(ner, X, Y)

i = 5
collect(zip(ner(X[i]), X[i], Y[i]))

i = 95
collect(zip(ner(X[i]), X[i], Y[i]))



StarcoderdataJupyterScriptsDedupFiltered