Update generate.py

generate.py

@@ -1,22 +1,28 @@
 import torch
-from torch.nn import functional as F
+import torch.nn.functional as F
+from evo_model import EvoDecoderModel
 
-def generate_text(model, tokenizer, prompt, max_length=100, temperature=1.0):
+def load_model(vocab_size, model_path="evo_decoder.pt", device="cpu"):
+    model = EvoDecoderModel(vocab_size)
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.to(device)
     model.eval()
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+    return model
 
-    generated = input_ids
-    memory = torch.zeros((1, input_ids.size(1), model.config.d_model)).to(model.device)
+def generate_text(model, tokenizer, prompt, max_new_tokens=50, temperature=1.0, device="cpu"):
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 
-    with torch.no_grad():
-        for _ in range(max_length):
-            outputs = model(generated, memory)
-            next_token_logits = outputs[:, -1, :] / temperature
-            probs = F.softmax(next_token_logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            generated = torch.cat((generated, next_token), dim=1)
+    for _ in range(max_new_tokens):
+        with torch.no_grad():
+            logits = model(input_ids)
+        next_token_logits = logits[:, -1, :] / temperature
+        probs = F.softmax(next_token_logits, dim=-1)
+        next_token_id = torch.multinomial(probs, num_samples=1)
 
-            if next_token.item() == tokenizer.eos_token_id:
-                break
+        input_ids = torch.cat([input_ids, next_token_id], dim=-1)
 
-    return tokenizer.decode(generated[0], skip_special_tokens=True)
+        if next_token_id.item() == tokenizer.eos_token_id:
+            break
+
+    output_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+    return output_text