Update README.md

README.md

@@ -1,3 +1,90 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+datasets:
+- Leore42/RETAN
+language:
+- en
+---
+
+```python
+This is an extremely tiny model that summarizes text into 4 words
+
+you will need to downlaod the config, tokenizer and model and use this pytho ncode as a starting point:
+
+import torch
+import tkinter as tk
+import json
+import torch.nn as nn
+import math
+
+class ThemeExtractor(nn.Module):
+    def __init__(self, vocab_size, d_model=64, nhead=4, num_layers=1, dropout=0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=128, dropout=dropout)
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.fc = nn.Linear(d_model, 1)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        emb = self.embedding(x) * math.sqrt(self.d_model)
+        emb = emb.transpose(0, 1)
+        enc_out = self.encoder(emb)
+        enc_out = enc_out.transpose(0, 1)
+        enc_out = self.dropout(enc_out)
+        logits = self.fc(enc_out).squeeze(-1)
+        return logits
+
+def load_model_and_tokenizer():
+    with open('config.json', 'r') as f:
+        config = json.load(f)
+    with open('tokenizer.json', 'r') as f:
+        vocab = json.load(f)
+    vocab_size = config["vocab_size"]
+    model = ThemeExtractor(vocab_size, d_model=64, nhead=4, num_layers=1, dropout=0.2)
+    model.load_state_dict(torch.load("theme_extractor.pth"))
+    return model, vocab, config
+
+def generate_text(model, vocab, config, input_text):
+    inv_vocab = {v: k for k, v in vocab.items()}
+    max_len = config["max_len"]
+    tokens = input_text.lower().split()
+    token_ids = [vocab.get(token, vocab["<unk>"]) for token in tokens]
+    if len(token_ids) < max_len:
+        token_ids += [vocab["<pad>"]] * (max_len - len(token_ids))
+    else:
+        token_ids = token_ids[:max_len]
+    input_tensor = torch.tensor([token_ids], dtype=torch.long).to(device)
+    model.eval()
+    with torch.no_grad():
+        logits = model(input_tensor)
+        probs = torch.sigmoid(logits).squeeze(0)
+    topk = torch.topk(probs, 4)
+    indices = topk.indices.cpu().numpy()
+    selected = sorted(indices, key=lambda i: i)
+    theme_words = [tokens[i] for i in selected if i < len(tokens)]
+    return ' '.join(theme_words)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model, vocab, config = load_model_and_tokenizer()
+
+def on_generate():
+    input_text = entry_input.get()
+    generated = generate_text(model, vocab, config, input_text)
+    label_output.config(text="Generated Themes: " + generated)
+
+root = tk.Tk()
+root.title("Theme Extractor")
+
+entry_input = tk.Entry(root, width=50)
+entry_input.pack(pady=10)
+
+button_generate = tk.Button(root, text="Generate Themes", command=on_generate)
+button_generate.pack(pady=10)
+
+label_output = tk.Label(root, text="Generated Themes: ", wraplength=400)
+label_output.pack(pady=10)
+
+root.mainloop()
+```