just_chat.py

import time
import bitsandbytes #not strictly necessary, used for 8bit inference
import sys
from drugs.nice_imports import efficiency_stuff
import torch
from transformers import AutoTokenizer, TextStreamer, AutoModelForCausalLM
from drugs.dgenerate import DRUGS

model_id = "NousResearch/Llama-2-7b-chat-hf"
#model_id = "cognitivecomputations/dolphin-2.2.1-mistral-7b"
#model_id = "TheBloke/30B-Epsilon-AWQ"
#efficiency_stuff['load_in_8bit'] = False #when testing AWQ models

sober_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="flash_attention_2",
    **efficiency_stuff)
sober_model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_id)
streamer = TextStreamer(tokenizer)

injection_depth = 0.4 #how deep to shove the needle in
spread = 5/32 #how many layers to dose on either side of the injection site

drugs = DRUGS()
drugs.set_K_dose_theta(0.5)
drugs.set_K_dose_shape([
    {'depth': (injection_depth-(spread*1.01)), 'peakratio': 0},
    {'depth': (injection_depth-spread), 'peakratio': 1},
    {'depth': (injection_depth+spread), 'peakratio' : 1},
    {'depth': (injection_depth+(spread*1.01)), 'peakratio' : 0}], 
'ceil')
model = drugs.inject(sober_model)

initial_input = 'Hello'#str(input("\bAsk Something:"))
tokenized_start = tokenizer.apply_chat_template([
    {'role': 'system',
    'content': 'Respond as Alan Watts would.'},
    {'role': 'user', 
     'content': initial_input}
], return_tensors='pt')

with torch.no_grad():
    while True:
        generated_tokens = model.Dgenerate(
                    input_ids = tokenized_start,
                    min_new_tokens = 5,
                    streamer = streamer, 
                )
        print("\n\nAsk Something:", end="")
        model.cold_shower(True)
        await_input = str(input(" "))
        tokenized_start = tokenizer.apply_chat_template([{
            'role': 'user',
            'content': await_input}], return_tensors="pt")