Update generate.py

generate.py

@@ -1,35 +1,22 @@
-# generate.py
 import torch
-from transformers import AutoTokenizer
-from evo_model import EvoDecoderModel
+from torch.nn import functional as F
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+def generate_text(model, tokenizer, prompt, max_length=100, temperature=1.0):
+    model.eval()
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-vocab_size = tokenizer.vocab_size
+    generated = input_ids
+    memory = torch.zeros((1, input_ids.size(1), model.config.d_model)).to(model.device)
 
-model = EvoDecoderModel(vocab_size=vocab_size).to(device)
-model.load_state_dict(torch.load("evo_decoder_model.pt", map_location=device))
-model.eval()
+    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)
 
-def generate_response(prompt, max_new_tokens=50, use_web=False):
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128)
-    input_ids = inputs["input_ids"].to(device)
+            if next_token.item() == tokenizer.eos_token_id:
+                break
 
-    for _ in range(max_new_tokens):
-        with torch.no_grad():
-            logits = model(input_ids)
-
-        next_token_logits = logits[:, -1, :]  # shape (B, vocab_size)
-        next_token_id = torch.argmax(next_token_logits, dim=-1).unsqueeze(0)  # shape (1, 1)
-
-        # Append to input
-        input_ids = torch.cat([input_ids, next_token_id], dim=1)
-
-        # Stop if EOS token
-        if next_token_id.item() in tokenizer.all_special_ids:
-            break
-
-    output_text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
-    return output_text[len(prompt):].strip()
+    return tokenizer.decode(generated[0], skip_special_tokens=True)