File size: 6,799 Bytes
930e8c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b65e3b1
930e8c0
 
 
 
 
 
 
 
b65e3b1
930e8c0
 
 
 
 
b65e3b1
930e8c0
b65e3b1
930e8c0
 
 
 
 
 
 
 
 
 
 
 
 
 
b65e3b1
930e8c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import os, xml.etree.ElementTree as ET, torch, torch.nn as nn, numpy as np, logging, requests
from collections import defaultdict
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, AutoModel
from sklearn.metrics.pairwise import cosine_similarity
from accelerate import Accelerator
from tqdm import tqdm

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class Config: E, H, N, C, B, M, S, V, W, L, D = 512, 32, 1024, 256, 128, 20000, 2048, 1e5, 4000, 2e-4, .15

class MyDataset(Dataset):
    def __init__(self, data, labels): self.data, self.labels = data, labels
    def __len__(self): return len(self.data)
    def __getitem__(self, index): return self.data[index], self.labels[index]

class MyModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MyModel, self).__init__()
        self.hidden, self.output = nn.Linear(input_size, hidden_size), nn.Linear(hidden_size, output_size)
        self.lstm, self.fc = nn.LSTM(input_size, hidden_size, batch_first=True), nn.Linear(hidden_size, output_size)
    def forward(self, x):
        x = torch.relu(self.hidden(x))
        h0, c0 = torch.zeros(1, x.size(0), hidden_size), torch.zeros(1, x.size(0), hidden_size)
        out, _ = self.lstm(x, (h0, c0))
        return self.fc(out[:, -1, :])

class MemoryNetwork:
    def __init__(self, memory_size, embedding_size):
        self.memory, self.usage = np.zeros((memory_size, embedding_size)), np.zeros(memory_size)
    def store(self, data):
        index = np.argmin(self.usage)
        self.memory[index], self.usage[index] = data, 1.0
    def retrieve(self, query):
        index = np.argmax(np.dot(self.memory, query))
        self.usage[index] += 1.0
        return self.memory[index]
    def update_usage(self): self.usage *= 0.99

class DM(nn.Module):
    def __init__(self, s):
        super(DM, self).__init__()
        self.s = nn.ModuleDict({sn: nn.ModuleList([self.cl(lp) for lp in l]) for sn, l in s.items()})
    def cl(self, lp):
        l = [nn.Linear(lp['input_size'], lp['output_size'])]
        if lp.get('batch_norm', True): l.append(nn.BatchNorm1d(lp['output_size']))
        a = lp.get('activation', 'relu')
        if a == 'relu': l.append(nn.ReLU(inplace=True))
        elif a == 'tanh': l.append(nn.Tanh())
        elif a == 'sigmoid': l.append(nn.Sigmoid())
        elif a == 'leaky_relu': l.append(nn.LeakyReLU(negative_slope=0.01, inplace=True))
        elif a == 'elu': l.append(nn.ELU(alpha=1.0, inplace=True))
        if dr := lp.get('dropout', 0.0): l.append(nn.Dropout(p=dr))
        return nn.Sequential(*l)
    def forward(self, x, sn=None):
        if sn: 
            for l in self.s[sn]: x = l(x)
        else:
            for sn, l in self.s.items():
                for l in l: x = l(x)
        return x

def parse_xml(file_path):
    t, r, l = ET.parse(file_path), ET.parse(file_path).getroot(), []
    for ly in r.findall('.//layer'):
        lp = {'input_size': int(ly.get('input_size', 128)), 'output_size': int(ly.get('output_size', 256)), 'activation': ly.get('activation', 'relu').lower()}
        l.append(lp)
    return l

def create_model_from_folder(folder_path):
    s = defaultdict(list)
    for r, d, f in os.walk(folder_path):
        for file in f:
            if file.endswith('.xml'):
                s[os.path.basename(r).replace('.', '_')].extend(parse_xml(os.path.join(r, file)))
    return DM(dict(s))

def create_embeddings_and_sentences(folder_path, model_name="pile-of-law/legalbert-large-1.7M-1", max_length=512):
    t, m, embeddings, ds = AutoTokenizer.from_pretrained(model_name), AutoModel.from_pretrained(model_name), [], []
    for r, d, f in os.walk(folder_path):
        for file in f:
            if file.endswith('.xml'):
                tree, root = ET.parse(os.path.join(r, file)), ET.parse(os.path.join(r, file)).getroot()
                for e in root.iter():
                    if e.text:
                        text = e.text.strip()
                        i = t(text, return_tensors="pt", truncation=True, padding=True, max_length=max_length)
                        with torch.no_grad():
                            embeddings.append(m(**i).last_hidden_state.mean(dim=1).numpy())
                        ds.append(text)
    return np.vstack(embeddings), ds

def query_vector_similarity(query, embeddings, ds, model_name="pile-of-law/legalbert-large-1.7M-2", max_length=512):
    t, m = AutoTokenizer.from_pretrained(model_name), AutoModel.from_pretrained(model_name)
    i = t(query, return_tensors="pt", truncation=True, padding=True, max_length=max_length)
    with torch.no_grad():
        qe = m(**i).last_hidden_state.mean(dim=1).numpy()
    return [ds[i] for i in cosine_similarity(qe, embeddings)[0].argsort()[-5:][::-1]]

def fetch_courtlistener_data(query):
    try:
        response = requests.get("https://nzlii.org/cgi-bin/sinosrch.cgi", params={"method": "auto", "query": query, "meta": "/nz", "results": "50", "format": "json"}, headers={"Accept": "application/json"}, timeout=10)
        response.raise_for_status()
        return [{"title": r.get("title", ""), "citation": r.get("citation", ""), "date": r.get("date", ""), "court": r.get("court", ""), "summary": r.get("summary", ""), "url": r.get("url", "")} for r in response.json().get("results", [])]
    except requests.exceptions.RequestException as e:
        logging.error(f"Failed to fetch data from NZLII API: {str(e)}")
        return []

def main():
    folder_path, model = 'data', create_model_from_folder('Xml_Data')
    logging.info(f"Created dynamic PyTorch model with sections: {list(model.s.keys())}")
    embeddings, ds = create_embeddings_and_sentences(folder_path)
    accelerator, optimizer, criterion, num_epochs = Accelerator(), torch.optim.Adam(model.parameters(), lr=0.001), nn.CrossEntropyLoss(), 10
    dataset, dataloader = MyDataset(torch.randn(1000, 10), torch.randint(0, 5, (1000,))), DataLoader(MyDataset(torch.randn(1000, 10), torch.randint(0, 5, (1000,))), batch_size=32, shuffle=True)
    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
    for epoch in range(num_epochs):
        model.train()
        for batch_data, batch_labels in dataloader:
            optimizer.zero_grad()
            outputs = model(batch_data)
            loss = criterion(outputs, batch_labels)
            accelerator.backward(loss)
            optimizer.step()
        logging.info(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
    query = "example query text"
    logging.info(f"Query results: {query_vector_similarity(query, embeddings, ds)}")
    logging.info(f"CourtListener API results: {fetch_courtlistener_data(query)}")

if __name__ == "__main__":
    main()