Update: refined tars_v1_model.py structure and logic

TARS-v1/tars_v1_model.py

@@ -1,49 +1,66 @@
+# tars_v1_model.py
+"""
+TARS-v1: A Hybrid Quantum-AI Model Combining BERT and GPT-Neo for Natural Language Understanding and Generation.
+
+SEO Optimized Tags: Quantum AI, GPT-Neo BERT Hybrid, HuggingFace Compatible, NLP Model, TARS-v1, AI Assistant, SubatomicError, PyTorch Transformers
+"""
+
 import torch
 import torch.nn as nn
-from transformers import BertModel, GPTNeoForCausalLM, AutoTokenizer
+from transformers import BertModel, GPTNeoForCausalLM, BertTokenizer, GPT2Tokenizer
 
 class TARSQuantumHybrid(nn.Module):
     """
-    TARSQuantumHybrid: A hybrid architecture combining BERT's semantic power 
-    with GPT-Neo's generative reasoning. Designed for advanced quantum-AI interaction.
-    
-    Compatible with Hugging Face model hub and deployable in transformers pipelines.
+    TARSQuantumHybrid: Combines the deep language understanding of BERT with the generative capabilities of GPT-Neo.
     """
-
     def __init__(self, bert_model="bert-base-uncased", gpt_model="EleutherAI/gpt-neo-125M"):
         super(TARSQuantumHybrid, self).__init__()
-        self.bert = BertModel.from_pretrained(bert_model)
-        self.gpt = GPTNeoForCausalLM.from_pretrained(gpt_model)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Load pretrained models
+        self.bert = BertModel.from_pretrained(bert_model).to(self.device)
+        self.gpt = GPTNeoForCausalLM.from_pretrained(gpt_model).to(self.device)
 
-        # Ensure compatibility across config keys
-        gpt_hidden_dim = getattr(self.gpt.config, "hidden_size", None) or getattr(self.gpt.config, "n_embd", 768)
-        self.embedding_proj = nn.Linear(self.bert.config.hidden_size, gpt_hidden_dim)
+        # BERT hidden size -> GPT input embedding projection
+        self.embedding_proj = nn.Linear(
+            self.bert.config.hidden_size,
+            self.gpt.config.hidden_size  # FIXED: correct attribute
+        ).to(self.device)
 
     def forward(self, input_ids, attention_mask=None, decoder_input_ids=None):
-        """
-        Forward pass through BERT + projected embeddings into GPT-Neo.
-
-        Parameters:
-        - input_ids: Token IDs for BERT
-        - attention_mask: Attention mask for BERT
-        - decoder_input_ids: Decoder inputs for GPT
-
-        Returns:
-        - GPTNeo output logits
-        """
-        # Step 1: Encode with BERT
+        # BERT processes input
         bert_output = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cls_embedding = bert_output.last_hidden_state[:, 0, :]  # [CLS] token embedding
+        cls_embedding = bert_output.last_hidden_state[:, 0, :]  # Extract [CLS] token
 
-        # Step 2: Project to GPT input dimension
+        # Project BERT CLS to GPT's input embedding space
         gpt_input = self.embedding_proj(cls_embedding).unsqueeze(1)
 
-        # Step 3: Feed into GPT
-        outputs = self.gpt(inputs_embeds=gpt_input, decoder_input_ids=decoder_input_ids)
-        return outputs
+        # GPT-Neo generates output based on embedded input
+        output = self.gpt(inputs_embeds=gpt_input, decoder_input_ids=decoder_input_ids)
+        return output
 
 if __name__ == "__main__":
-    # Test load and save
+    # Load tokenizers for test
+    bert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+    gpt_tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
+
+    # Initialize model
     model = TARSQuantumHybrid()
+    model.eval()
+
+    # Test Input
+    sample_text = "What is quantum consciousness?"
+    tokens = bert_tokenizer(sample_text, return_tensors="pt").to(model.device)
+
+    # Dummy decoder input (for GPT)
+    decoder_input_ids = torch.tensor([[gpt_tokenizer.bos_token_id]]).to(model.device)
+
+    # Forward pass
+    with torch.no_grad():
+        output = model(input_ids=tokens['input_ids'],
+                       attention_mask=tokens['attention_mask'],
+                       decoder_input_ids=decoder_input_ids)
+    
+    # Save model weights
     torch.save(model.state_dict(), "tars_v1.pt")
     print("✅ TARS-v1 model saved successfully as 'tars_v1.pt'")