File size: 3,003 Bytes
db7bc8a cd7d193 db7bc8a cd7d193 db7bc8a cd7d193 db7bc8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
---
license: apache-2.0
datasets:
- Leore42/RETAN
language:
- en
---
This is an extremely tiny model that summarizes text into 4 words
you will need to download the config, tokenizer and model and use this pytho ncode as a starting point:
```python
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()
``` |