File size: 1,345 Bytes
1c98e83 b28676e 100e3bb b37c655 f718bd4 b37c655 f718bd4 b37c655 ccae0a9 100e3bb 1c98e83 3d17dd0 ef57a4b 3d17dd0 b28676e ccae0a9 5259900 b28676e f718bd4 b37c655 1c98e83 f718bd4 b28676e f718bd4 b28676e 7dc4300 b37c655 |
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 |
import torch
import torch.nn.functional as F
from transformers import GPT2Tokenizer
from evo_decoder import EvoDecoder
from search_utils import web_search
# 🔧 Load model and tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
model = EvoDecoder(
vocab_size=tokenizer.vocab_size,
d_model=256,
nhead=4,
num_layers=3,
dim_feedforward=512
).to(device)
model.load_state_dict(torch.load("evo_decoder.pt", map_location=device))
model.eval()
@torch.no_grad()
def generate_response(question, context="", use_rag=False, temperature=1.0):
if not context and use_rag:
context = web_search(question)
prompt = f"Context: {context}\nQuestion: {question}\nAnswer:"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
for _ in range(128):
logits = model(input_ids)
logits = logits[:, -1, :] / temperature
probs = F.softmax(logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
input_ids = torch.cat((input_ids, next_token), dim=1)
if next_token.item() == tokenizer.eos_token_id:
break
output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
return output[len(prompt):].strip()
|