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 # Fixed: Reduced from 10240 to reasonable value
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:
# Fixed: Use device_map="auto" but ensure proper tensor handling
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,
)
# Get the actual device of the model
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()
# Disabled model compilation as it can cause issues with quantized models
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)
# Fixed: More reasonable target length calculation
target_length = MAX_CONTEXT_LENGTH - MAX_NEW_TOKENS - 100 # Safety buffer
print(f"[Current prompt tokens: {len(tokens)}, Target: {target_length}]")
if len(tokens) > target_length:
print(f"[Truncating context: {len(tokens)} -> ~{target_length} tokens]")
# Calculate available space for conversation
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 # Safety buffer
print(f"[System tokens: {system_tokens}, Input tokens: {current_input_tokens}, Available for history: {available_tokens}]")
if available_tokens <= 100: # Need minimum space for meaningful history
# If no space for history, just use system prompt + current input
print("[Using minimal context - no conversation history]")
return f"{self.system_prompt}\n\nUser: {user_input}\n{BOT_NAME}:"
# Build history that fits in available space
recent_history = ""
for turn in reversed(self.memory[:-1]): # Start from most recent, excluding current
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
# Construct final prompt
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()
# Stop if we hit user indicators
if any(line_lower.startswith(indicator) for indicator in user_indicators):
break
# Keep the line if it's not empty and doesn't contain problematic phrases
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()
# Don't return empty responses due to over-aggressive cleaning
if not result and response.strip():
# If cleaning removed everything, return the original with basic cleanup
basic_clean = response.strip()
# Just remove obvious user indicators
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]")
# Fixed: Proper input preparation with device handling
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}]")
# Move inputs to the correct device
if hasattr(self.model, 'device'):
device = self.model.device
else:
device = next(self.model.parameters()).device
print(f"[Moving inputs to device: {device}]")
# Handle multi-device models (quantized models might spread across devices)
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}")
# For quantized models, just ensure input_ids are on the right device
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, # Fixed: Use the corrected value
"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 ['.', '!', '?'])
# Early stopping conditions - be less aggressive
if len(full_response.strip()) > 30: # Reduced from 50
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
# Less aggressive stopping - allow longer responses
if word_count >= 200 and sentence_count >= 4: # Increased thresholds
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)
# Fixed: Better handling of mixed precision for quantized models
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:
# For quantized models, don't use autocast as it can interfere
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()
- Downloads last month
- 6
Inference Providers
NEW
This model isn't deployed by any Inference Provider.
🙋
Ask for provider support