90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
|
||
|
|
import sys
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description="Quintus Interactive Chat")
|
||
|
|
parser.add_argument("--model_path", type=str, default="iamrahulreddy/Quintus", help="Model repo ID or local weights directory")
|
||
|
|
parser.add_argument("--trust_remote_code", action="store_true", help="Allow custom code from the model repository.")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
model_path = args.model_path
|
||
|
|
print(f"Loading Quintus from {model_path}...")
|
||
|
|
try:
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=args.trust_remote_code)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
model_path,
|
||
|
|
device_map="auto",
|
||
|
|
dtype=torch.float16,
|
||
|
|
trust_remote_code=args.trust_remote_code
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error loading model: {e}")
|
||
|
|
print(f"Ensure '{model_path}' exists and contains the model weights.")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
# Defining stopping criteria
|
||
|
|
stop_tokens = ["<|endoftext|>", "<|im_end|>"]
|
||
|
|
eos_token_ids = [tokenizer.eos_token_id] if tokenizer.eos_token_id is not None else []
|
||
|
|
for token in stop_tokens:
|
||
|
|
t_id = tokenizer.convert_tokens_to_ids(token)
|
||
|
|
if t_id is not None and t_id not in eos_token_ids:
|
||
|
|
eos_token_ids.append(t_id)
|
||
|
|
|
||
|
|
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||
|
|
|
||
|
|
conversation_history = [
|
||
|
|
{"role": "system", "content": "You are Quintus, a highly capable AI assistant created by Muskula Rahul. You are helpful, precise, and logically sound."}
|
||
|
|
]
|
||
|
|
|
||
|
|
print()
|
||
|
|
print("Quintus Chat (type 'quit' to exit)")
|
||
|
|
print()
|
||
|
|
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
user_input = input("You: ").strip()
|
||
|
|
if user_input.lower() in ["quit", "exit"]:
|
||
|
|
print("\nGoodbye!")
|
||
|
|
break
|
||
|
|
if not user_input:
|
||
|
|
continue
|
||
|
|
|
||
|
|
conversation_history.append({"role": "user", "content": user_input})
|
||
|
|
|
||
|
|
prompt = tokenizer.apply_chat_template(
|
||
|
|
conversation_history,
|
||
|
|
tokenize=False,
|
||
|
|
add_generation_prompt=True
|
||
|
|
)
|
||
|
|
|
||
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
||
|
|
|
||
|
|
print("Quintus: ", end="", flush=True)
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
outputs = model.generate(
|
||
|
|
**inputs,
|
||
|
|
max_new_tokens=512,
|
||
|
|
temperature=0.7,
|
||
|
|
top_p=0.9,
|
||
|
|
do_sample=True,
|
||
|
|
streamer=streamer,
|
||
|
|
pad_token_id=tokenizer.eos_token_id,
|
||
|
|
eos_token_id=eos_token_ids
|
||
|
|
)
|
||
|
|
|
||
|
|
# Extract response for history
|
||
|
|
generated_ids = outputs[0][inputs.input_ids.shape[-1]:]
|
||
|
|
assistant_response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
|
||
|
|
conversation_history.append({"role": "assistant", "content": assistant_response})
|
||
|
|
print()
|
||
|
|
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
print("\n\nGoodbye!")
|
||
|
|
break
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|