import sys
import subprocess
print("Installing Unsloth (Fast-Track Mode)... Please wait ~30-40 seconds.")
subprocess.run([sys.executable, "-m", "pip", "install", "--quiet",
"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git",
"xformers", "trl", "peft", "transformers", "accelerate", "bitsandbytes"])
import site
import importlib
importlib.invalidate_caches()
site.main()
print("Installation complete! Loading UCortex Model...")
try:
from unsloth import FastLanguageModel
except ModuleNotFoundError:
import os
sys.path.append('/usr/local/lib/python3.12/dist-packages')
from unsloth import FastLanguageModel
import torch
from transformers import TextIteratorStreamer
from threading import Thread
model_name = "Uzaib52/ucortex-llama3-8b"
max_seq_length = 2048
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
load_in_4bit = True,
device_map = "cuda"
)
FastLanguageModel.for_inference(model)
print("UCortex Chatbot is active! Type 'exit' to stop.\n" + "="*50)
prompt_format = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
while True:
user_input = input("\n You: ")
if user_input.lower() in ["exit", "quit"]:
print(" UCortex: Bye!")
break
if not user_input.strip():
continue
formatted_prompt = prompt_format.format(user_input)
inputs = tokenizer([formatted_prompt], return_tensors="pt").to("cuda")
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=512, use_cache=True)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
print("UCortex: ", end="", flush=True)
for new_text in streamer:
print(new_text, end="", flush=True)
print()