Update x.py

x.py

@@ -5,24 +5,35 @@ import xml.etree.ElementTree as ET
 import torch
 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 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
 
 # Initialize colorama
 init(autoreset=True)
+logging.basicConfig(level=logging.INFO, format='\033[92m%(asctime)s - %(levelname)s - %(message)s\033[0m')
 
-# Set file path and output path
 file_path = 'data/'
 output_path = 'output/'
 
 # Create output path if it doesn't exist
 if not os.path.exists(output_path):
-    os.makedirs(output_path)
-    os.chmod(output_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)  # Set full r/w permissions
+    try:
+        os.makedirs(output_path)
+        os.chmod(output_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)  # Set full r/w permissions
+    except PermissionError:
+        print(f"Permission denied: '{output_path}'")
+        # Handle the error or try a different path
 
 # Ensure necessary files are created with full r/w permissions
 def ensure_file(file_path):
@@ -31,124 +42,101 @@ def ensure_file(file_path):
             pass
         os.chmod(file_path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)  # Set full r/w permissions
 
-# Define a simple memory augmentation module
+class MagicStateLayer(nn.Module):
+    def __init__(self, size):
+        super().__init__()
+        self.state = nn.Parameter(torch.randn(size))
+
+    def forward(self, x):
+        return x + self.state
+
 class MemoryAugmentationLayer(nn.Module):
-    def __init__(self, size: int):
-        super(MemoryAugmentationLayer, self).__init__()
+    def __init__(self, size):
+        super().__init__()
         self.memory = nn.Parameter(torch.randn(size))
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x):
         return x + self.memory
 
 class HybridAttentionLayer(nn.Module):
-    def __init__(self, size: int):
-        super(HybridAttentionLayer, self).__init__()
+    def __init__(self, size):
+        super().__init__()
         self.attention = nn.MultiheadAttention(size, num_heads=8)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x.unsqueeze(1)  # Add sequence dimension
+    def forward(self, x):
+        x = x.unsqueeze(1)
         attn_output, _ = self.attention(x, x, x)
         return attn_output.squeeze(1)
 
 class DynamicFlashAttentionLayer(nn.Module):
-    def __init__(self, size: int):
-        super(DynamicFlashAttentionLayer, self).__init__()
+    def __init__(self, size):
+        super().__init__()
         self.attention = nn.MultiheadAttention(size, num_heads=8)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x.unsqueeze(1)  # Add sequence dimension
+    def forward(self, x):
+        x = x.unsqueeze(1)
         attn_output, _ = self.attention(x, x, x)
         return attn_output.squeeze(1)
 
 class DynamicModel(nn.Module):
     def __init__(self, sections: Dict[str, List[Dict[str, Any]]]):
-        super(DynamicModel, self).__init__()
-        self.sections = nn.ModuleDict()
-        if not sections:
-            sections = {
-                'default': [{
-                    'input_size': 128,
-                    'output_size': 256,
-                    'activation': 'relu',
-                    'batch_norm': True,
-                    'dropout': 0.1
-                }]
-            }
-        for section_name, layers in sections.items():
-            self.sections[section_name] = nn.ModuleList()
-            for layer_params in layers:
-                print(f"Creating layer in section '{section_name}' with params: {layer_params}")
-                self.sections[section_name].append(self.create_layer(layer_params))
-
-    def create_layer(self, layer_params: Dict[str, Any]) -> nn.Module:
-        layers = []
-        layers.append(nn.Linear(layer_params['input_size'], layer_params['output_size']))
-        if layer_params.get('batch_norm', False):
-            layers.append(nn.BatchNorm1d(layer_params['output_size']))
-        activation = layer_params.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))
-        elif activation is not None:
-            raise ValueError(f"Unsupported activation function: {activation}")
-        if dropout_rate := layer_params.get('dropout', 0.0):
-            layers.append(nn.Dropout(p=dropout_rate))
-        if hidden_layers := layer_params.get('hidden_layers', []):
-            for hidden_layer_params in hidden_layers:
-                layers.append(self.create_layer(hidden_layer_params))
-        if layer_params.get('memory_augmentation', True):
-            layers.append(MemoryAugmentationLayer(layer_params['output_size']))
-        if layer_params.get('hybrid_attention', True):
-            layers.append(HybridAttentionLayer(layer_params['output_size']))
-        if layer_params.get('dynamic_flash_attention', True):
-            layers.append(DynamicFlashAttentionLayer(layer_params['output_size']))
-        return nn.Sequential(*layers)
-
-    def forward(self, x: torch.Tensor, section_name: Optional[str] = None) -> torch.Tensor:
-        if section_name is not None:
-            if section_name not in self.sections:
-                raise KeyError(f"Section '{section_name}' not found in model")
-            for layer in self.sections[section_name]:
-                x = layer(x)
+        super().__init__()
+        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'])]
+        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))
+        if lp.get('memory_augmentation', False):
+            l.append(MemoryAugmentationLayer(lp['output_size']))
+        if lp.get('hybrid_attention', False):
+            l.append(HybridAttentionLayer(lp['output_size']))
+        if lp.get('dynamic_flash_attention', False):
+            l.append(DynamicFlashAttentionLayer(lp['output_size']))
+        if lp.get('magic_state', False):
+            l.append(MagicStateLayer(lp['output_size']))
+        return nn.Sequential(*l)
+
+    def forward(self, x, sn=None):
+        if sn:
+            for l in self.sections[sn]:
+                x = l(x)
         else:
-            for section_name, layers in self.sections.items():
-                for layer in layers:
-                    x = layer(x)
+            for sn, layers in self.sections.items():
+                for l in layers:
+                    x = l(x)
         return x
 
-def parse_xml_file(file_path: str) -> List[Dict[str, Any]]:
-    tree = ET.parse(file_path)
-    root = tree.getroot()
-    layers = []
+def parse_xml_file(file_path):
+    tree, root, layers = ET.parse(file_path), ET.parse(file_path).getroot(), []
     for layer in root.findall('.//layer'):
-        layer_params = {}
-        layer_params['input_size'] = int(layer.get('input_size', 128))
-        layer_params['output_size'] = int(layer.get('output_size', 256))
-        layer_params['activation'] = layer.get('activation', 'relu').lower()
-        if layer_params['activation'] not in ['relu', 'tanh', 'sigmoid', 'none']:
-            raise ValueError(f"Unsupported activation function: {layer_params['activation']}")
-        if layer_params['input_size'] <= 0 or layer_params['output_size'] <= 0:
+        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:
             raise ValueError("Layer dimensions must be positive integers")
-        layers.append(layer_params)
+        layers.append(lp)
     if not layers:
-        layers.append({
-            'input_size': 128,
-            'output_size': 256,
-            'activation': 'relu'
-        })
+        layers.append({'input_size': 128, 'output_size': 256, 'activation': 'relu'})
     return layers
 
-def create_model_from_folder(folder_path: str) -> DynamicModel:
+def create_model_from_folder(folder_path):
     sections = defaultdict(list)
     if not os.path.exists(folder_path):
-        print(f"Warning: Folder {folder_path} does not exist. Creating model with default configuration.")
+        logging.warning(f"Folder {folder_path} does not exist. Creating model with default configuration.")
         return DynamicModel({})
     xml_files_found = False
     for root, dirs, files in os.walk(folder_path):
@@ -157,68 +145,152 @@ def create_model_from_folder(folder_path: str) -> DynamicModel:
                 xml_files_found = True
                 file_path = os.path.join(root, file)
                 try:
-                    layers = parse_xml_file(file_path)
-                    section_name = os.path.basename(root).replace('.', '_')
-                    sections[section_name].extend(layers)
+                    sections[os.path.basename(root).replace('.', '_')].extend(parse_xml_file(file_path))
                 except Exception as e:
-                    print(f"Error processing {file_path}: {str(e)}")
+                    logging.error(f"Error processing {file_path}: {str(e)}")
     if not xml_files_found:
-        print("Warning: No XML files found. Creating model with default configuration.")
+        logging.warning("No XML files found. Creating model with default configuration.")
         return DynamicModel({})
     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), []
+    for root, dirs, files in os.walk(folder_path):
+        for file in files:
+            if file.endswith('.xml'):
+                file_path = os.path.join(root, file)
+                try:
+                    tree, root = ET.parse(file_path), ET.parse(file_path).getroot()
+                    for elem in root.iter():
+                        if elem.text:
+                            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)
+                            doc_store.append(text)
+                except Exception as e:
+                    logging.error(f"Error processing {file_path}: {str(e)}")
+    return vector_store, 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)
+    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]]
+
+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"}
+    try:
+        response = requests.get(base_url, params=params, 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 []
+
+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.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)
+
+    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 training_step(self, data, labels, optimizer, criterion):
+        optimizer.zero_grad()
+        outputs = self.forward(data)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        return loss.item()
+
+    def validation_step(self, data, labels, criterion):
+        with torch.no_grad():
+            outputs = self.forward(data)
+            loss = criterion(outputs, labels)
+        return loss.item()
+
+    def predict(self, input):
+        self.eval()
+        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():
-    print(Fore.CYAN + "Starting conversion...")
-
-    # Create the dynamic model from the folder
-    model = create_model_from_folder(file_path)
-    print(f"Created dynamic PyTorch model with sections: {list(model.sections.keys())}")
-
-    # Print the model architecture
-    print(model)
-
-    # Ensure the input tensor size matches the expected input size
-    first_section = next(iter(model.sections.keys()))
-    first_layer = model.sections[first_section][0]
-    input_features = first_layer[0].in_features
-    sample_input = torch.randn(1, input_features)
-    output = model(sample_input)
-    print(f"Sample output shape: {output.shape}")
-
-    # Training setup
-    accelerator = Accelerator()
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
-    criterion = nn.CrossEntropyLoss()
-    num_epochs = 10
-    dataset = TensorDataset(
-        torch.randn(100, input_features),
-        torch.randint(0, 2, (100,))
-    )
-    train_dataloader = DataLoader(
-        dataset,
-        batch_size=8,  # Reduced batch size
-        shuffle=True
-    )
-    model, optimizer, train_dataloader = accelerator.prepare(
-        model, optimizer, train_dataloader
-    )
-    scaler = GradScaler()  # Mixed precision training
-
-    # Training loop
+    folder_path, model = 'data', create_model_from_folder('data')
+    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)
+    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
     for epoch in range(num_epochs):
         model.train()
         total_loss = 0
-        for batch_idx, (inputs, labels) in enumerate(train_dataloader):
+        for batch_data, batch_labels in dataloader:
             optimizer.zero_grad()
-            with autocast():  # Mixed precision training
-                outputs = model(inputs)
-                loss = criterion(outputs, labels)
-            scaler.scale(loss).backward()
-            scaler.step(optimizer)
-            scaler.update()
+            outputs = model(batch_data)
+            loss = criterion(outputs, batch_labels)
+            accelerator.backward(loss)
+            optimizer.step()
             total_loss += loss.item()
-        avg_loss = total_loss / len(train_dataloader)
-        print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}")
+        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)
+    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()