File size: 3,072 Bytes
373f237 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
# chat.py
import os
import gc
import torch
from transformers import LlamaTokenizer, LlamaForCausalLM
# =============================
# Configuration
# =============================
MODEL_PATH = r"C:\Users\JAY\Downloads\Chatdoc\ChatDoctor\pretrained"
MAX_NEW_TOKENS = 200
TEMPERATURE = 0.5
TOP_K = 50
REPETITION_PENALTY = 1.1
# Detect device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading model from {MODEL_PATH} on {device}...")
# =============================
# Load Tokenizer and Model
# =============================
tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)
model = LlamaForCausalLM.from_pretrained(
MODEL_PATH,
device_map="auto", # automatically dispatch weights to GPU
torch_dtype=torch.float16, # half precision for faster inference
low_cpu_mem_usage=True # optimize CPU memory
)
# DO NOT call model.to(device) when using device_map="auto"
generator = model.generate
print("✅ Model loaded successfully!\n")
# =============================
# Chat History
# =============================
history = ["ChatDoctor: I am ChatDoctor, what medical questions do you have?"]
# =============================
# Response Function
# =============================
def get_response(user_input):
global history
human_invitation = "Patient: "
doctor_invitation = "ChatDoctor: "
# Append user input
history.append(human_invitation + user_input)
# Build prompt
prompt = "\n".join(history) + "\n" + doctor_invitation
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
# Generate response
with torch.no_grad():
output_ids = generator(
input_ids,
max_new_tokens=MAX_NEW_TOKENS,
do_sample=True,
temperature=TEMPERATURE,
top_k=TOP_K,
repetition_penalty=REPETITION_PENALTY
)
# Decode response
full_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
response = full_output[len(prompt):].strip()
# Clean if the model repeats the patient prompt
if response.startswith("Patient:"):
response = response[len("Patient:"):].strip()
# Append model response to history
history.append(doctor_invitation + response)
# Free memory
del input_ids, output_ids
gc.collect()
torch.cuda.empty_cache()
return response
# =============================
# CLI Chat
# =============================
if __name__ == "__main__":
print("\n=== ChatDoctor is ready! Type your questions. ===\n")
while True:
try:
user_input = input("Patient: ").strip()
if user_input.lower() in ["exit", "quit"]:
print("Exiting ChatDoctor. Goodbye!")
break
response = get_response(user_input)
print("ChatDoctor: " + response + "\n")
except KeyboardInterrupt:
print("\nExiting ChatDoctor. Goodbye!")
break
|