Update x.py

x.py

@@ -1,5 +1,4 @@
 import os
-import glob
 import stat
 import xml.etree.ElementTree as ET
 import torch
@@ -7,17 +6,14 @@ import torch.nn as nn
 import torch.nn.functional as F
 import logging
 import requests
-import faiss
 from collections import defaultdict
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any
 from colorama import Fore, Style, init
 from accelerate import Accelerator
 from torch.utils.data import DataLoader, TensorDataset
-from torch.cuda.amp import GradScaler, autocast
 from transformers import AutoTokenizer, AutoModel
 from sentence_transformers import SentenceTransformer
-from sentence_transformers.readers import ParagraphReader
-from sentence_transformers.uniformer import Uniformer
+import numpy as np
 
 # Initialize colorama
 init(autoreset=True)
@@ -84,46 +80,50 @@ class DynamicModel(nn.Module):
         self.sections = nn.ModuleDict({sn: nn.ModuleList([self.create_layer(lp) for lp in layers]) for sn, layers in sections.items()})
 
     def create_layer(self, lp):
-        l = [nn.Linear(lp['input_size'], lp['output_size'])]
+        layers = [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))
+            layers.append(nn.BatchNorm1d(lp['output_size']))
+        activation = lp.get('activation', 'relu')
+        if activation == 'relu':
+            layers.append(nn.ReLU(inplace=True))
+        elif activation == 'tanh':
+            layers.append(nn.Tanh())
+        elif activation == 'sigmoid':
+            layers.append(nn.Sigmoid())
+        elif activation == 'leaky_relu':
+            layers.append(nn.LeakyReLU(negative_slope=0.01, inplace=True))
+        elif activation == 'elu':
+            layers.append(nn.ELU(alpha=1.0, inplace=True))
+        if dropout := lp.get('dropout', 0.0):
+            layers.append(nn.Dropout(p=dropout))
         if lp.get('memory_augmentation', False):
-            l.append(MemoryAugmentationLayer(lp['output_size']))
+            layers.append(MemoryAugmentationLayer(lp['output_size']))
         if lp.get('hybrid_attention', False):
-            l.append(HybridAttentionLayer(lp['output_size']))
+            layers.append(HybridAttentionLayer(lp['output_size']))
         if lp.get('dynamic_flash_attention', False):
-            l.append(DynamicFlashAttentionLayer(lp['output_size']))
+            layers.append(DynamicFlashAttentionLayer(lp['output_size']))
         if lp.get('magic_state', False):
-            l.append(MagicStateLayer(lp['output_size']))
-        return nn.Sequential(*l)
+            layers.append(MagicStateLayer(lp['output_size']))
+        return nn.Sequential(*layers)
 
-    def forward(self, x, sn=None):
-        if sn:
-            for l in self.sections[sn]:
-                x = l(x)
+    def forward(self, x, section_name=None):
+        if section_name:
+            for layer in self.sections[section_name]:
+                x = layer(x)
         else:
-            for sn, layers in self.sections.items():
-                for l in layers:
-                    x = l(x)
+            for section_name, layers in self.sections.items():
+                for layer in layers:
+                    x = layer(x)
         return x
 
 def parse_xml_file(file_path):
     tree, root, layers = ET.parse(file_path), ET.parse(file_path).getroot(), []
     for layer in root.findall('.//layer'):
-        lp = {'input_size': int(layer.get('input_size', 128)), 'output_size': int(layer.get('output_size', 256)), 'activation': layer.get('activation', 'relu').lower()}
+        lp = {
+            'input_size': int(layer.get('input_size', 128)),
+            'output_size': int(layer.get('output_size', 256)),
+            'activation': layer.get('activation', 'relu').lower()
+        }
         if lp['activation'] not in ['relu', 'tanh', 'sigmoid', 'none']:
             raise ValueError(f"Unsupported activation function: {lp['activation']}")
         if lp['input_size'] <= 0 or lp['output_size'] <= 0:
@@ -154,7 +154,10 @@ def create_model_from_folder(folder_path):
     return DynamicModel(dict(sections))
 
 def create_embeddings_and_stores(folder_path, model_name="sentence-transformers/all-MiniLM-L6-v2"):
-    tokenizer, model, vector_store, doc_store = AutoTokenizer.from_pretrained(model_name), AutoModel.from_pretrained(model_name), faiss.IndexFlatL2(384), []
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModel.from_pretrained(model_name)
+    doc_store = []
+    embeddings_list = []
     for root, dirs, files in os.walk(folder_path):
         for file in files:
             if file.endswith('.xml'):
@@ -166,23 +169,26 @@ def create_embeddings_and_stores(folder_path, model_name="sentence-transformers/
                             text = elem.text.strip()
                             inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
                             with torch.no_grad():
-                                embeddings = model(**inputs).last_hidden_state.mean(dim=1).numpy()
-                            vector_store.add(embeddings)
+                                embeddings = model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()
+                            embeddings_list.append(embeddings)
                             doc_store.append(text)
                 except Exception as e:
                     logging.error(f"Error processing {file_path}: {str(e)}")
-    return vector_store, doc_store
+    return embeddings_list, doc_store
 
-def query_vector_store(query, vector_store, doc_store, model_name="sentence-transformers/all-MiniLM-L6-v2"):
-    tokenizer, model = AutoTokenizer.from_pretrained(model_name), AutoModel.from_pretrained(model_name)
+def query_embeddings(query, embeddings_list, doc_store, model_name="sentence-transformers/all-MiniLM-L6-v2"):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModel.from_pretrained(model_name)
     inputs = tokenizer(query, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
-        query_embedding = model(**inputs).last_hidden_state.mean(dim=1).numpy()
-    D, I = vector_store.search(query_embedding, k=5)
-    return [doc_store[i] for i in I[0]]
+        query_embedding = model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()
+    similarities = [np.dot(query_embedding, emb.T) for emb in embeddings_list]
+    top_k_indices = np.argsort(similarities, axis=0)[-5:][::-1]
+    return [doc_store[i] for i in top_k_indices]
 
 def fetch_courtlistener_data(query):
-    base_url, params = "https://nzlii.org/cgi-bin/sinosrch.cgi", {"method": "auto", "query": query, "meta": "/nz", "results": "50", "format": "json"}
+    base_url = "https://nzlii.org/cgi-bin/sinosrch.cgi"
+    params = {"method": "auto", "query": query, "meta": "/nz", "results": "50", "format": "json"}
     try:
         response = requests.get(base_url, params=params, headers={"Accept": "application/json"}, timeout=10)
         response.raise_for_status()
@@ -194,14 +200,14 @@ def fetch_courtlistener_data(query):
 class CustomModel(nn.Module):
     def __init__(self, model_name="distilbert-base-uncased"):
         super().__init__()
-        self.model_name = model_name
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.encoder = AutoModel.from_pretrained(model_name)
         self.hidden_size = self.encoder.config.hidden_size
-        self.dropout = nn.Dropout(p=0.2)
-        self.fc1 = nn.Linear(self.hidden_size, 64)
-        self.fc2 = nn.Linear(64, 32)
-        self.fc3 = nn.Linear(32, 16)
+        self.dropout = nn.Dropout(p=0.3)
+        self.fc1 = nn.Linear(self.hidden_size, 128)
+        self.fc2 = nn.Linear(128, 64)
+        self.fc3 = nn.Linear(64, 32)
+        self.fc4 = nn.Linear(32, 16)
         self.memory = nn.LSTM(self.hidden_size, 64, bidirectional=True, batch_first=True)
         self.memory_fc1 = nn.Linear(64 * 2, 32)
         self.memory_fc2 = nn.Linear(32, 16)
@@ -212,7 +218,8 @@ class CustomModel(nn.Module):
         x = outputs.last_hidden_state.mean(dim=1)
         x = self.dropout(F.relu(self.fc1(x)))
         x = self.dropout(F.relu(self.fc2(x)))
-        x = self.fc3(x)
+        x = self.dropout(F.relu(self.fc3(x)))
+        x = self.fc4(x)
         return x
 
     def training_step(self, data, labels, optimizer, criterion):
@@ -234,45 +241,17 @@ class CustomModel(nn.Module):
         with torch.no_grad():
             return self.forward(input)
 
-class CustomModelInference(nn.Module):
-    def __init__(self, model_name="distilbert-base-uncased"):
-        super().__init__()
-        self.model_name = model_name
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.encoder = AutoModel.from_pretrained(model_name)
-        self.hidden_size = self.encoder.config.hidden_size
-        self.dropout = nn.Dropout(p=0.2)
-        self.fc1 = nn.Linear(self.hidden_size, 64)
-        self.fc2 = nn.Linear(64, 32)
-        self.fc3 = nn.Linear(32, 16)
-        self.reader = ParagraphReader("data/docstore.json")
-        self.model_embedding = SentenceTransformer('sentence-transformers/multilingual-v2')
-        self.vectorstore = Uniformer("distilusembert-base-nli-mean-tokens", torch.nn.CrossEntropyLoss(), margin=0.5, temperature=0.1, top_k=4)
-
-    def forward(self, data):
-        tokens = self.tokenizer(data, return_tensors="pt", truncation=True, padding=True)
-        outputs = self.encoder(**tokens)
-        x = outputs.last_hidden_state.mean(dim=1)
-        x = self.dropout(F.relu(self.fc1(x)))
-        x = self.dropout(F.relu(self.fc2(x)))
-        x = self.fc3(x)
-        return x
-
-    def infer(self, input):
-        self.eval()
-        with torch.no_grad():
-            return self.forward(input)
-
-    def update_memory(self, data):
-        embeddings = self.model_embedding.encode(data, convert_to_tensor=True)
-        self.vectorstore.add(embeddings)
-
 def main():
-    folder_path, model = 'data', create_model_from_folder('data')
+    folder_path = 'data'
+    model = create_model_from_folder(folder_path)
     logging.info(f"Created dynamic PyTorch model with sections: {list(model.sections.keys())}")
-    vector_store, doc_store = create_embeddings_and_stores(folder_path)
-    accelerator, optimizer, criterion, num_epochs = Accelerator(), torch.optim.Adam(model.parameters(), lr=0.001), nn.CrossEntropyLoss(), 10
-    dataset, dataloader = TensorDataset(torch.randn(100, 128), torch.randint(0, 2, (100,))), DataLoader(TensorDataset(torch.randn(100, 128), torch.randint(0, 2, (100,))), batch_size=16, shuffle=True)
+    embeddings_list, doc_store = create_embeddings_and_stores(folder_path)
+    accelerator = Accelerator()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+    criterion = nn.CrossEntropyLoss()
+    num_epochs = 10
+    dataset = TensorDataset(torch.randn(100, 128), torch.randint(0, 2, (100,)))
+    dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
     model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
     for epoch in range(num_epochs):
         model.train()
@@ -287,10 +266,10 @@ def main():
         avg_loss = total_loss / len(dataloader)
         logging.info(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}")
     query = "example query text"
-    results = query_vector_store(query, vector_store, doc_store)
+    results = query_embeddings(query, embeddings_list, doc_store)
     logging.info(f"Query results: {results}")
     courtlistener_data = fetch_courtlistener_data(query)
     logging.info(f"CourtListener API results: {courtlistener_data}")
 
 if __name__ == "__main__":
-    main()
+    main()