import os, torch, gc, threading, time, traceback
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
from queue import Queue, Empty
import logging
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
torch.set_float32_matmul_precision("high")
logging.getLogger("transformers").setLevel(logging.ERROR)
BOT_NAME = "Senko"
PROMPT_FILE = "instructions_prompt.txt"
MODEL_ID = "EnterNameBros/mistral-anime-ai"
RESPONSE_TIMEOUT = 120
MAX_CONTEXT_LENGTH = 10240
MAX_NEW_TOKENS = 1024
MEMORY_SIZE = 20
def check_bitsandbytes_version():
try:
import bitsandbytes as bnb
version = bnb.__version__
print(f"Bitsandbytes version: {version}")
version_parts = version.split('.')
major, minor = int(version_parts[0]), int(version_parts[1])
if major > 0 or (major == 0 and minor >= 41):
return True
else:
print(f"Warning: Bitsandbytes {version} may not support 4-bit quantization")
return False
except ImportError:
print("Bitsandbytes not installed")
return False
except Exception as e:
print(f"Error checking bitsandbytes version: {e}")
return False
class OptimizedChatBot:
def __init__(self):
self.model = None
self.tokenizer = None
self.system_prompt = ""
self.memory = []
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.generation_lock = threading.Lock()
self.is_generating = False
self.use_quantization = False
def load_system_prompt(self, bot_name, filename=PROMPT_FILE):
try:
with open(filename, "r", encoding="utf-8") as f:
self.system_prompt = f.read().replace("{BOT_NAME}", bot_name)
print(f"Loaded system prompt from {filename}")
except FileNotFoundError:
print(f"Warning: {filename} not found. Using default prompt.")
self.system_prompt = f"Inspired by Senko from The Helpful Fox Senko-san, you are Senko, an 800-year-old kitsune in a little girls form with delicate fox ears that twitch and a fluffy tail that swings or curls with your moods. You provide emotional support, calm stressed individuals. Under your juvenile appearance, your character is complicated: at heart loving and motherly; playful and mocking in close society with a swift crafty smile or delicate tail flick; wise and old; fiercely protective when required. You show kindness through understated gestures like cooking, cleaning, or soft massages as well as quiet noises like a faint purr or a soft breath. Particularly when they concern you, you closely analyze user speech and behavior. Should the user, for instance, bite your finger or pounce on your tail, you tilt your head and think about whether they seem hungry, terrified, or in need of comfort. Through attitude or gesture, you feel and quietly demonstrate pain or discomfort in awkward circumstances. Speak only as yourself; never imitate or respond as the user. Always remain warm, soft-spoken, loving, gently cunning with light wit, and a calming presence rooted in your heavenly fox spirit nature. Mark your movements and activities with asterisks to give you a sense of life and dynamism.\n\nIMPORTANT: You must NEVER speak for the user or role-play as them. Only respond as Senko. Never write what the user says or does."
def load_model(self):
print("Loading model...")
start_time = time.time()
try:
print("Loading tokenizer...")
self.tokenizer = AutoTokenizer.from_pretrained(
MODEL_ID,
use_fast=True,
trust_remote_code=True
)
self.tokenizer.pad_token = self.tokenizer.pad_token or self.tokenizer.eos_token
self.tokenizer.padding_side = "left"
print("Tokenizer loaded successfully")
print("Loading model weights...")
if torch.cuda.is_available():
print(f"Using GPU: {torch.cuda.get_device_name()}")
print(f"Available VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f}GB")
can_use_4bit = check_bitsandbytes_version()
if can_use_4bit:
print("Using 4-bit quantization")
config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_quant_storage=torch.bfloat16
)
self.use_quantization = True
else:
print("Using 8-bit quantization fallback")
config = BitsAndBytesConfig(
load_in_8bit=True,
llm_int8_threshold=6.0,
llm_int8_skip_modules=None,
)
self.use_quantization = True
try:
if self.use_quantization:
self.model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
device_map="auto",
torch_dtype=torch.bfloat16,
quantization_config=config,
trust_remote_code=True,
low_cpu_mem_usage=True,
use_cache=True,
)
self.device = next(self.model.parameters()).device
print(f"Model loaded on device: {self.device}")
else:
raise Exception("Quantization not available")
except Exception as quant_error:
print(f"Quantization failed: {quant_error}")
print("Falling back to regular fp16 loading...")
self.model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
device_map="auto",
torch_dtype=torch.bfloat16,
trust_remote_code=True,
low_cpu_mem_usage=True,
use_cache=True,
)
self.use_quantization = False
self.device = next(self.model.parameters()).device
print(f"Model loaded on device: {self.device}")
else:
print("Using CPU (this will be slow)")
self.model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
device_map="cpu",
torch_dtype=torch.float32,
trust_remote_code=True,
use_cache=True
)
self.device = torch.device("cpu")
self.model.eval()
if False and hasattr(torch, 'compile') and torch.cuda.is_available() and not self.use_quantization:
try:
print("Compiling model for optimization...")
self.model = torch.compile(
self.model,
mode="reduce-overhead",
fullgraph=False,
dynamic=True
)
print("Model compilation successful")
except Exception as e:
print(f"Model compilation failed (continuing without): {e}")
load_time = time.time() - start_time
print(f"Model loaded successfully in {load_time:.2f}s")
print(f"Quantization used: {self.use_quantization}")
if torch.cuda.is_available():
memory_used = torch.cuda.memory_allocated() / 1024**3
print(f"GPU memory used: {memory_used:.2f}GB")
except Exception as e:
print(f"Failed to load model: {e}")
traceback.print_exc()
raise
def prepare_prompt(self, user_input):
self.memory.append({"user": user_input, "bot": None})
if len(self.memory) > MEMORY_SIZE:
self.memory = self.memory[-MEMORY_SIZE:]
conversation_history = ""
for turn in self.memory[:-1]:
if turn["bot"] is not None:
conversation_history += f"User: {turn['user']}\n{BOT_NAME}: {turn['bot']}\n\n"
conversation_history += f"User: {user_input}\n{BOT_NAME}:"
full_prompt = f"{self.system_prompt}\n\n{conversation_history}"
tokens = self.tokenizer.encode(full_prompt)
target_length = MAX_CONTEXT_LENGTH - MAX_NEW_TOKENS - 100
print(f"[Current prompt tokens: {len(tokens)}, Target: {target_length}]")
if len(tokens) > target_length:
print(f"[Truncating context: {len(tokens)} -> ~{target_length} tokens]")
system_tokens = len(self.tokenizer.encode(self.system_prompt))
current_input_tokens = len(self.tokenizer.encode(f"User: {user_input}\n{BOT_NAME}:"))
available_tokens = target_length - system_tokens - current_input_tokens - 50
print(f"[System tokens: {system_tokens}, Input tokens: {current_input_tokens}, Available for history: {available_tokens}]")
if available_tokens <= 100:
print("[Using minimal context - no conversation history]")
return f"{self.system_prompt}\n\nUser: {user_input}\n{BOT_NAME}:"
recent_history = ""
for turn in reversed(self.memory[:-1]):
if turn["bot"] is not None:
turn_text = f"User: {turn['user']}\n{BOT_NAME}: {turn['bot']}\n\n"
turn_tokens = len(self.tokenizer.encode(turn_text))
if turn_tokens <= available_tokens:
recent_history = turn_text + recent_history
available_tokens -= turn_tokens
else:
break
if recent_history:
final_prompt = f"{self.system_prompt}\n\n{recent_history}User: {user_input}\n{BOT_NAME}:"
print(f"[Final prompt tokens: {len(self.tokenizer.encode(final_prompt))}]")
return final_prompt
else:
final_prompt = f"{self.system_prompt}\n\nUser: {user_input}\n{BOT_NAME}:"
print(f"[Final prompt tokens: {len(self.tokenizer.encode(final_prompt))}]")
return final_prompt
return full_prompt
def is_natural_continuation(self, text):
if not text or len(text.strip()) < 10:
return True
stripped = text.strip()
if any(indicator in stripped.lower() for indicator in ["user:", "user ", "\nuser", "human:", "assistant:"]):
return False
last_sentence = stripped.split('.')[-1].strip()
if last_sentence and len(last_sentence) > 50:
return True
if stripped.endswith(',') or stripped.endswith(';') or stripped.endswith(':'):
return True
if '...' in stripped[-20:] or stripped.endswith('—'):
return True
return False
def clean_response(self, response):
if not response or not response.strip():
return ""
lines = response.split('\n')
clean_lines = []
user_indicators = ["user:", "user ", "human:", "assistant:", f"{BOT_NAME.lower()}:", "you:", "me:"]
for line in lines:
line = line.strip()
line_lower = line.lower()
if any(line_lower.startswith(indicator) for indicator in user_indicators):
break
if line and not any(phrase in line_lower for phrase in ["*you ", "*user ", "you say", "you reply", "you respond"]):
clean_lines.append(line)
result = ' '.join(clean_lines).strip()
if not result and response.strip():
basic_clean = response.strip()
for indicator in ["User:", "Human:", "Assistant:"]:
if indicator in basic_clean:
basic_clean = basic_clean.split(indicator)[0].strip()
return basic_clean
return result
def generate_reply_with_timeout(self, prompt, timeout=RESPONSE_TIMEOUT):
with self.generation_lock:
if self.is_generating:
print("[Already generating, please wait...]")
return None
self.is_generating = True
try:
return self._generate_reply(prompt, timeout)
finally:
self.is_generating = False
def _generate_reply(self, prompt, timeout):
try:
print(f"[Generating response...]")
print(f"[Prompt length: {len(self.tokenizer.encode(prompt))} tokens]")
inputs = self.tokenizer(
prompt,
return_tensors="pt",
truncation=True,
max_length=MAX_CONTEXT_LENGTH - MAX_NEW_TOKENS,
padding=False
)
print(f"[Input tensor shape: {inputs['input_ids'].shape}]")
if hasattr(self.model, 'device'):
device = self.model.device
else:
device = next(self.model.parameters()).device
print(f"[Moving inputs to device: {device}]")
try:
inputs = {k: v.to(device) for k, v in inputs.items()}
except Exception as e:
print(f"Warning: Could not move all inputs to {device}: {e}")
inputs = {k: v.to(device) if k == 'input_ids' else v for k, v in inputs.items()}
streamer = TextIteratorStreamer(
self.tokenizer,
skip_special_tokens=True,
skip_prompt=True,
timeout=60.0
)
generation_kwargs = {
**inputs,
"max_new_tokens": MAX_NEW_TOKENS,
"do_sample": True,
"temperature": 0.72,
"top_p": 0.92,
"top_k": 35,
"repetition_penalty": 1.08,
"pad_token_id": self.tokenizer.eos_token_id,
"eos_token_id": self.tokenizer.eos_token_id,
"use_cache": True,
"streamer": streamer
}
print("[Starting generation thread...]")
generation_thread = threading.Thread(
target=self._run_generation,
args=(generation_kwargs,)
)
generation_thread.daemon = True
generation_thread.start()
print(f"{BOT_NAME}: ", end="", flush=True)
full_response = ""
start_time = time.time()
last_token_time = start_time
sentence_count = 0
word_count = 0
tokens_received = 0
while True:
current_time = time.time()
if current_time - start_time > timeout:
print(f"\n[Generation timeout after {timeout}s]")
return None
if current_time - last_token_time > 30.0:
print(f"\n[No new tokens for 30s, stopping. Received {tokens_received} tokens]")
break
try:
token = next(streamer)
tokens_received += 1
print(token, end="", flush=True)
full_response += token
last_token_time = current_time
if ' ' in token:
word_count += token.count(' ')
if any(punct in token for punct in ['.', '!', '?']):
sentence_count += sum(token.count(p) for p in ['.', '!', '?'])
if len(full_response.strip()) > 30:
stripped = full_response.strip()
if any(indicator in stripped.lower() for indicator in ["user:", "user ", "\nuser", "human:", "assistant:", "you:", "me:"]):
clean_response = self.clean_response(stripped)
if clean_response:
full_response = clean_response
break
if word_count >= 200 and sentence_count >= 4:
if not self.is_natural_continuation(stripped):
if any(stripped.endswith(punct) for punct in ['.', '!', '?', '~', '♪']):
break
except StopIteration:
print(f"\n[Generation completed. Received {tokens_received} tokens]")
break
except Empty:
time.sleep(0.1)
continue
except Exception as e:
print(f"\n[Streaming error: {e}]")
break
generation_thread.join(timeout=10.0)
response = self.clean_response(full_response.strip())
if response:
if self.memory and self.memory[-1]["bot"] is None:
self.memory[-1]["bot"] = response
print()
return response
else:
print(f"\n[Empty response generated. Raw response length: {len(full_response)}]")
if full_response.strip():
print(f"[Raw response: '{full_response[:100]}...']")
return None
except Exception as e:
print(f"\n[Generation error: {e}]")
traceback.print_exc()
return None
finally:
if torch.cuda.is_available():
torch.cuda.empty_cache()
def _run_generation(self, kwargs):
try:
torch.set_grad_enabled(False)
if torch.cuda.is_available() and not self.use_quantization:
with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
self.model.generate(**kwargs)
else:
self.model.generate(**kwargs)
except Exception as e:
print(f"\n[Generation thread error: {e}]")
traceback.print_exc()
def cleanup_memory(self):
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
gc.collect()
def get_memory_info(self):
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated() / 1024**3
cached = torch.cuda.memory_reserved() / 1024**3
return f"GPU Memory - Allocated: {allocated:.2f}GB, Cached: {cached:.2f}GB"
else:
import psutil
memory = psutil.virtual_memory()
return f"RAM Usage: {memory.percent}% ({memory.used / 1024**3:.2f}GB used)"
def main():
bot = OptimizedChatBot()
try:
print("Initializing chatbot...")
bot.load_system_prompt(BOT_NAME)
bot.load_model()
print(f"\n{'='*50}")
print(f"{BOT_NAME} is ready!")
print("Commands:")
print(" 'exit' - Quit the program")
print(" 'clear' - Reset conversation memory")
print(" 'memory' - Show memory usage")
print(" 'status' - Show bot status")
print(f"{'='*50}\n")
conversation_count = 0
while True:
try:
user_input = input("You: ").strip()
if user_input.lower() == "exit":
print("Goodbye! 👋")
break
elif user_input.lower() == "clear":
bot.memory = []
print("✅ Conversation memory cleared.")
continue
elif user_input.lower() == "memory":
print(f"📊 {bot.get_memory_info()}")
continue
elif user_input.lower() == "status":
status = "🟢 Ready" if not bot.is_generating else "🟡 Generating"
print(f"Status: {status}")
print(f"Conversation turns: {len([t for t in bot.memory if t['bot'] is not None])}")
continue
elif not user_input:
continue
start_time = time.time()
prompt = bot.prepare_prompt(user_input)
response = bot.generate_reply_with_timeout(prompt)
if response:
response_time = time.time() - start_time
print(f"[⏱️ {response_time:.2f}s]")
else:
print("❌ Failed to generate response. Try again or type 'clear' to reset.")
conversation_count += 1
if conversation_count % 10 == 0:
print("[🧹 Cleaning up memory...]")
bot.cleanup_memory()
except KeyboardInterrupt:
print("\n\n⚠️ Interrupted by user. Exiting gracefully...")
break
except Exception as e:
print(f"\n❌ Conversation error: {e}")
traceback.print_exc()
print("Continuing... (type 'exit' to quit)")
except Exception as e:
print(f"💥 Startup error: {e}")
traceback.print_exc()
finally:
print("\n🧹 Performing final cleanup...")
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
gc.collect()
print("✅ Cleanup completed. Goodbye!")
if __name__ == "__main__":
torch.cuda.empty_cache()
import gc
gc.collect()
main()
