File size: 2,672 Bytes
38b4eff | 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 | import sys
with open('/home/workspace/Projects/NeuralAI/services/neural_core_service.py', 'r') as f:
content = f.read()
# Replace generate_response_stream signature and logic
old_stream = '''def generate_response_stream(prompt, max_tokens=256, temperature=0.7, system_prompt=None):
global model, tokenizer, inference_count
if model is None or tokenizer is None:
yield "Model not loaded."
return
if system_prompt is None:
system_prompt = DEFAULT_SYSTEM_PROMPT
try:
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.append({"role": "user", "content": prompt})
if tokenizer.chat_template:
full = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
else:
full = f"system\\n{system_prompt}\\nuser\\n{prompt}" if system_prompt else f"user\\n{prompt}"'''
new_stream = '''def generate_response_stream(prompt, max_tokens=256, temperature=0.7, system_prompt=None, history=None):
global model, tokenizer, inference_count
if model is None or tokenizer is None:
yield "Model not loaded."
return
if system_prompt is None:
system_prompt = DEFAULT_SYSTEM_PROMPT
try:
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
if history:
messages.extend(history)
else:
messages.append({"role": "user", "content": prompt})
if tokenizer.chat_template:
full = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
else:
full = ""
for msg in messages:
full += f"<|im_start|>{msg['role']}\\n{msg['content']}<|im_end|>\\n"
full += "<|im_start|>assistant\\n"'''
if old_stream not in content:
print("Failed to find old stream logic!")
sys.exit(1)
content = content.replace(old_stream, new_stream)
# Replace the chat endpoint call
old_call = 'for chunk in generate_response_stream(prompt, max_tokens, temperature, system_prompt=system_prompt):'
new_call = 'for chunk in generate_response_stream(prompt, max_tokens, temperature, system_prompt=system_prompt, history=messages):'
if old_call not in content:
print("Failed to find chat endpoint call!")
sys.exit(1)
content = content.replace(old_call, new_call)
with open('/home/workspace/Projects/NeuralAI/services/neural_core_service.py', 'w') as f:
f.write(content)
print("Successfully patched.")
|