reverse-llm-chatbot.py

import requests
import json
import sys
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")

SYSTEM_PROMPT = "You are a helpful assistant."
messages = [{"role": "system", "content": SYSTEM_PROMPT}]

def inverse_chat_history(messages):
    prompt = tokenizer.apply_chat_template(messages+[{"role": "user", "content": "<<CLIP>>"}], tokenize=False, add_generation_prompt=True)
    clip_idx = prompt.find('<<CLIP>>')
    prompt = prompt[:clip_idx]
    return prompt

# llama.cpp server configuration
SERVER_URL = "http://10.0.0.168:8084/completion"
HEADERS = {"Content-Type": "application/json"}

def stream_response(prompt):
    data = {
        "prompt": prompt,
        "n_predict": 1000,
        "stream": True
    }
    completion = ''
    try:
        with requests.post(SERVER_URL, headers=HEADERS, json=data, stream=True) as response:
            response.raise_for_status()
            
            for line in response.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith('data: '):
                        json_data = json.loads(decoded_line[6:])
                        if 'content' in json_data:
                            sys.stdout.write(json_data['content'])
                            completion += json_data['content'] 
                            sys.stdout.flush()
                        if 'stop' in json_data and json_data['stop']:
                            break
    except requests.exceptions.RequestException as e:
        print(f"Error: {e}")
        
    return completion

def main():
    print("Welcome to the Reverse LLM Chatbot!")
    print("Type 'exit' to quit the application.")
    
    while True:
        print("\nUser: ", end='')
        user_input = stream_response(inverse_chat_history(messages))
        messages.append({"role": "user", "content": user_input})
        
        ai_input = input("\n\"AI\": ")
        if ai_input.lower() == 'exit':
            print("Goodbye!")
            break

        messages.append({"role": "assistant", "content": ai_input})

if __name__ == "__main__":
    main()