Update generate.py

generate.py

@@ -24,6 +24,10 @@ def generate_response(prompt, use_web=False, max_length=100, top_k=40):
     input_text = context + f"User: {prompt}\nAssistant:"
     input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
 
+    # Truncate to avoid exceeding model's positional encoding limit
+    if input_ids.size(1) > 512:
+        input_ids = input_ids[:, -512:]
+
     for _ in range(max_length):
         with torch.no_grad():
             logits = model(input_ids)