Create app.py

app.py

@@ -0,0 +1,336 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+import math
+from transformers import PreTrainedTokenizerFast
+import numpy as np
+from typing import Optional, List, Dict
+
+class NeuromodulatedAttention(nn.Module):
+    def __init__(self, d_model: int, num_heads: int):
+        super().__init__()
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        
+        self.qkv = nn.Linear(d_model, 3 * d_model)
+        self.out_proj = nn.Linear(d_model, d_model)
+        
+        # Neuromodulation
+        self.dopamine_gate = nn.Linear(d_model, num_heads)
+        self.serotonin_gate = nn.Linear(d_model, num_heads)
+        self.memory_decay = nn.Parameter(torch.ones(num_heads) * 0.99)
+        self.forget_gate = nn.Linear(d_model, num_heads)
+        self.attention_mask = nn.Parameter(torch.ones(num_heads))
+        
+        # Memory
+        self.register_buffer('memory_state', torch.zeros(1, num_heads, 1, self.head_dim))
+
+    def update_memory(self, new_info: torch.Tensor, dopamine: torch.Tensor, forget: torch.Tensor):
+        self.memory_state = (
+            self.memory_state * self.memory_decay.view(1, -1, 1, 1) * 
+            (1 - forget.unsqueeze(-1)) + 
+            dopamine.unsqueeze(-1) * new_info
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_length, _ = x.shape
+        
+        # Neuromodulators
+        dopamine = torch.sigmoid(self.dopamine_gate(x.mean(dim=1)))
+        serotonin = torch.sigmoid(self.serotonin_gate(x.mean(dim=1)))
+        forget = torch.sigmoid(self.forget_gate(x.mean(dim=1)))
+        
+        # Attention computation
+        qkv = self.qkv(x)
+        qkv = qkv.reshape(batch_size, seq_length, 3, self.num_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        
+        # Include memory
+        k = torch.cat([k, self.memory_state.expand(batch_size, -1, -1, -1)], dim=2)
+        v = torch.cat([v, self.memory_state.expand(batch_size, -1, -1, -1)], dim=2)
+        
+        # Attention with neuromodulation
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        scores = scores * serotonin.view(batch_size, self.num_heads, 1, 1)
+        scores = scores * self.attention_mask.view(1, -1, 1, 1)
+        
+        attention = F.softmax(scores, dim=-1)
+        x = torch.matmul(attention, v)
+        
+        # Update memory
+        self.update_memory(x.mean(dim=2), dopamine, forget)
+        
+        x = x.transpose(1, 2).reshape(batch_size, seq_length, self.d_model)
+        return self.out_proj(x)
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
+        super().__init__()
+        self.attention = NeuromodulatedAttention(d_model, num_heads)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.ReLU(),
+            nn.Linear(d_ff, d_model),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attention(self.norm1(x))
+        x = x + self.ff(self.norm2(x))
+        return x
+
+class NeuroTransformer(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int = 256,
+        num_heads: int = 4,
+        num_layers: int = 3,
+        d_ff: int = 512,
+        dropout: float = 0.1,
+        max_seq_length: int = 128
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_encoding = self._create_positional_encoding(max_seq_length, d_model)
+        self.layers = nn.ModuleList([
+            TransformerBlock(d_model, num_heads, d_ff, dropout)
+            for _ in range(num_layers)
+        ])
+        self.final_layer = nn.Linear(d_model, vocab_size)
+        self.dropout = nn.Dropout(dropout)
+
+    def _create_positional_encoding(self, max_seq_length: int, d_model: int) -> torch.Tensor:
+        pos_encoding = torch.zeros(max_seq_length, d_model)
+        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pos_encoding[:, 0::2] = torch.sin(position * div_term)
+        pos_encoding[:, 1::2] = torch.cos(position * div_term)
+        return pos_encoding.unsqueeze(0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.embedding(x) * math.sqrt(self.d_model)
+        x = x + self.pos_encoding[:, :x.size(1)].to(x.device)
+        x = self.dropout(x)
+        
+        for layer in self.layers:
+            x = layer(x)
+            
+        return self.final_layer(x)
+
+    def generate(
+        self,
+        tokenizer: PreTrainedTokenizerFast,
+        prompt: str,
+        max_length: int = 100,
+        temperature: float = 0.7,
+        top_k: int = 50,
+        top_p: float = 0.9
+    ) -> str:
+        self.eval()
+        input_ids = tokenizer.encode(prompt, return_tensors="pt")
+        
+        with torch.no_grad():
+            for _ in range(max_length):
+                outputs = self(input_ids)
+                next_token_logits = outputs[:, -1, :] / temperature
+                
+                # Top-k
+                if top_k > 0:
+                    indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                    next_token_logits[indices_to_remove] = float('-inf')
+                
+                # Top-p
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    next_token_logits[indices_to_remove] = float('-inf')
+                
+                probs = F.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+                
+                if next_token.item() == tokenizer.eos_token_id:
+                    break
+                    
+                input_ids = torch.cat([input_ids, next_token], dim=1)
+        
+        return tokenizer.decode(input_ids[0], skip_special_tokens=True)
+
+class TextGenerator:
+    def __init__(self):
+        self.tokenizer = PreTrainedTokenizerFast.from_pretrained('gpt2')
+        self.model = NeuroTransformer(vocab_size=self.tokenizer.vocab_size)
+        
+    def train_on_text(
+        self,
+        text: str,
+        epochs: int,
+        learning_rate: float,
+        batch_size: int,
+        progress=gr.Progress()
+    ) -> str:
+        encodings = self.tokenizer(text, truncation=True, padding=True, return_tensors="pt")
+        input_ids = encodings['input_ids']
+        
+        dataset = torch.utils.data.TensorDataset(input_ids, input_ids)
+        dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
+        
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=learning_rate)
+        criterion = nn.CrossEntropyLoss()
+        
+        logs = []
+        self.model.train()
+        
+        for epoch in progress.tqdm(range(epochs)):
+            total_loss = 0
+            for batch in dataloader:
+                optimizer.zero_grad()
+                input_ids, labels = batch
+                outputs = self.model(input_ids)
+                loss = criterion(outputs.view(-1, outputs.size(-1)), labels.view(-1))
+                loss.backward()
+                optimizer.step()
+                total_loss += loss.item()
+            
+            avg_loss = total_loss / len(dataloader)
+            logs.append(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")
+        
+        return "\n".join(logs)
+    
+    def generate(
+        self,
+        prompt: str,
+        max_length: int,
+        temperature: float,
+        top_k: int,
+        top_p: float
+    ) -> str:
+        return self.model.generate(
+            self.tokenizer,
+            prompt,
+            max_length=max_length,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p
+        )
+
+# Create Gradio interface
+generator = TextGenerator()
+
+demo = gr.Blocks()
+
+with demo:
+    gr.Markdown("# Neuromodulated Text Generator")
+    
+    with gr.Tab("Train"):
+        with gr.Row():
+            with gr.Column():
+                train_input = gr.Textbox(
+                    label="Training Text",
+                    placeholder="Enter text to train on...",
+                    lines=5
+                )
+                train_button = gr.Button("Train Model")
+                
+            with gr.Column():
+                epochs_slider = gr.Slider(
+                    label="Epochs",
+                    minimum=1,
+                    maximum=50,
+                    value=10,
+                    step=1
+                )
+                lr_slider = gr.Slider(
+                    label="Learning Rate",
+                    minimum=1e-5,
+                    maximum=1e-3,
+                    value=1e-4,
+                    step=1e-5
+                )
+                batch_slider = gr.Slider(
+                    label="Batch Size",
+                    minimum=1,
+                    maximum=32,
+                    value=4,
+                    step=1
+                )
+                
+        train_output = gr.Textbox(label="Training Log")
+        
+    with gr.Tab("Generate"):
+        with gr.Row():
+            with gr.Column():
+                prompt_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter text prompt...",
+                    lines=2
+                )
+                generate_button = gr.Button("Generate Text")
+                
+            with gr.Column():
+                length_slider = gr.Slider(
+                    label="Max Length",
+                    minimum=10,
+                    maximum=500,
+                    value=100,
+                    step=10
+                )
+                temp_slider = gr.Slider(
+                    label="Temperature",
+                    minimum=0.1,
+                    maximum=2.0,
+                    value=0.7,
+                    step=0.1
+                )
+                topk_slider = gr.Slider(
+                    label="Top-k",
+                    minimum=0,
+                    maximum=100,
+                    value=50,
+                    step=1
+                )
+                topp_slider = gr.Slider(
+                    label="Top-p",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.05
+                )
+                
+        generate_output = gr.Textbox(label="Generated Text")
+
+    train_button.click(
+        fn=generator.train_on_text,
+        inputs=[
+            train_input,
+            epochs_slider,
+            lr_slider,
+            batch_slider
+        ],
+        outputs=train_output
+    )
+
+    generate_button.click(
+        fn=generator.generate,
+        inputs=[
+            prompt_input,
+            length_slider,
+            temp_slider,
+            topk_slider,
+            topp_slider
+        ],
+        outputs=generate_output
+    )
+
+if __name__ == "__main__":
+    demo.launch()