Lotta work here

include/lm/models/attention.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include "lm/core/tensor.hpp"
+#include <vector>
+#include <memory>
+
+namespace lm {
+
+class MultiHeadAttention {
+public:
+    MultiHeadAttention(size_t d_model, size_t num_heads, float dropout = 0.1f);
+    
+    std::vector<Tensor> parameters() const;
+    void set_training(bool training);
+    Tensor forward(const Tensor& query, const Tensor& key, const Tensor& value, 
+                   const Tensor& mask = Tensor()) const;
+    
+private:
+    Tensor split_heads(const Tensor& x) const;
+    Tensor combine_heads(const Tensor& x) const;
+    Tensor scaled_dot_product_attention(const Tensor& q, const Tensor& k, 
+                                        const Tensor& v, const Tensor& mask) const;
+    Tensor apply_dropout(const Tensor& input, float dropout_rate) const;
+    
+    size_t d_model_;
+    size_t num_heads_;
+    size_t d_k_;
+    float dropout_;
+    bool training_ = false;
+    
+    Tensor w_q_;
+    Tensor w_k_;
+    Tensor w_v_;
+    Tensor w_o_;
+};
+
+} // namespace lm

include/lm/models/feed_forward.hpp

@@ -0,0 +1,32 @@
+#pragma once
+
+#include "lm/core/tensor.hpp"
+#include <vector>
+
+namespace lm {
+
+class FeedForward {
+public:
+    FeedForward(size_t d_model, size_t d_ff, float dropout = 0.1f);
+    
+    std::vector<Tensor> parameters() const;
+    void set_training(bool training);
+    Tensor forward(const Tensor& input) const;
+    
+private:
+    Tensor apply_dropout(const Tensor& input, float dropout_rate) const;
+    Tensor gelu(const Tensor& input) const;
+    
+    size_t d_model_;
+    size_t d_ff_;
+    float dropout_;
+    bool training_ = false;
+    
+    Tensor w1_;
+    Tensor b1_;
+    Tensor w2_;
+    Tensor b2_;
+};
+
+} // namespace lm
+

include/lm/models/language_model.hpp

@@ -0,0 +1,45 @@
+// lm/models/language_model.hpp
+#pragma once
+
+#include "../core/tensor.hpp"
+#include "../tokenizer/bpe_tokenizer.hpp"
+#include <vector>
+
+namespace lm {
+
+class LanguageModel {
+public:
+    LanguageModel(size_t vocab_size, size_t embedding_dim, size_t hidden_dim, size_t num_layers);
+    
+    Tensor forward(const Tensor& input);
+    Tensor operator()(const Tensor& input) { return forward(input); }
+
+    void save(const std::string& path) const;
+    void load(const std::string& path);
+    
+    // Parameter access methods
+    std::vector<Tensor> parameters() const;
+    std::unordered_map<std::string, Tensor> named_parameters() const;
+    void set_parameter(const std::string& name, const Tensor& param);
+    
+    void train();
+    void eval();
+    
+private:
+    size_t vocab_size_, embedding_dim_, hidden_dim_, num_layers_;
+    
+    // Model parameters
+    Tensor embedding_weight_;
+    Tensor lstm_weight_ih_;
+    Tensor lstm_weight_hh_;
+    Tensor lstm_bias_ih_;
+    Tensor lstm_bias_hh_;
+    Tensor output_weight_;
+    Tensor output_bias_;
+    
+    bool is_training_;
+    std::unordered_map<std::string, Tensor> parameters_;
+
+};
+
+} // namespace lm

include/lm/models/layer_norm.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "lm/core/tensor.hpp"
+#include <vector>
+
+namespace lm {
+
+class LayerNorm {
+public:
+    LayerNorm(size_t d_model, float eps = 1e-5f);
+    
+    std::vector<Tensor> parameters() const;
+    void set_training(/*bool training*/);
+    Tensor forward(const Tensor& input) const;
+    
+private:
+    size_t d_model_;
+    float eps_;
+    
+    Tensor gamma_;
+    Tensor beta_;
+};
+
+} // namespace lm

include/lm/models/transformer.hpp

@@ -0,0 +1,34 @@
+#pragma once
+
+#include "lm/core/tensor.hpp"
+#include "lm/models/transformer_block.hpp"
+#include <vector>
+#include <memory>
+#include <cmath>
+
+namespace lm {
+
+class Transformer {
+public:
+    Transformer(size_t vocab_size, size_t d_model, size_t num_heads, 
+                size_t d_ff, size_t num_layers, size_t max_seq_len, float dropout = 0.1f);
+    
+    std::vector<Tensor> parameters() const;
+    void set_training(bool training);
+    Tensor forward(const Tensor& input, const Tensor& mask);
+    Tensor forward(const Tensor& input);
+    
+private:
+    Tensor apply_dropout(const Tensor& input, float dropout_rate);
+    
+    size_t vocab_size_, d_model_, num_heads_, d_ff_, num_layers_, max_seq_len_;
+    float dropout_;
+    bool training_ = false;
+    
+    Tensor embedding_;
+    Tensor positional_encoding_;
+    Tensor output_layer_;
+    std::vector<std::unique_ptr<TransformerBlock>> transformer_blocks_;
+};
+
+} // namespace lm

include/lm/models/transformer_block.hpp

@@ -0,0 +1,32 @@
+#pragma once
+
+#include "lm/core/tensor.hpp"
+#include "lm/models/attention.hpp"
+#include "lm/models/feed_forward.hpp"
+#include "lm/models/layer_norm.hpp"
+#include <memory>
+#include <vector>
+
+namespace lm {
+
+class TransformerBlock {
+public:
+    TransformerBlock(size_t d_model, size_t num_heads, size_t d_ff, float dropout);
+    
+    std::vector<Tensor> parameters() const;
+    void set_training(bool training);
+    Tensor forward(const Tensor& input, const Tensor& mask = Tensor()) const;
+    
+private:
+    size_t d_model_, num_heads_, d_ff_;
+    float dropout_;
+    bool training_ = false;
+    
+    std::unique_ptr<MultiHeadAttention> attention_;
+    std::unique_ptr<FeedForward> feed_forward_;
+    std::unique_ptr<LayerNorm> norm1_;
+    std::unique_ptr<LayerNorm> norm2_;
+};
+
+} // namespace lm
+