hf_chat.py

# hf_chat.py

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM
)

print("=== Hugging Face Chat ===")

MODEL_ID = input("Hugging Face model ID: ").strip()

print("""
Prompt mode:
1) Raw text
2) Chat template
""")

mode = input("Select mode: ").strip()

device = "cuda" if torch.cuda.is_available() else "cpu"

print("Using device:", device)

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token


print("Loading model...")

dtype = torch.float16 if device == "cuda" else torch.float32

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=dtype,
    device_map="auto",
    trust_remote_code=True
)

model.eval()

print("\nLoaded!")
print("Type Ctrl+C to exit.\n")


history = []


while True:
    try:
        user = input("\nUser: ")

        if mode == "1":
            prompt = input(
                "Full prompt format (leave empty for raw input): "
            )

            if prompt.strip():
                prompt = prompt
            else:
                prompt = user


        elif mode == "2":

            history.append({
                "role": "user",
                "content": user
            })

            prompt = tokenizer.apply_chat_template(
                history,
                tokenize=False,
                add_generation_prompt=True
            )

        else:
            print("Invalid mode")
            break


        inputs = tokenizer(
            prompt,
            return_tensors="pt"
        ).to(model.device)


        with torch.no_grad():

            output = model.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.4,
                top_p=0.9,
                do_sample=True,
                repetition_penalty=1.05,
                pad_token_id=tokenizer.eos_token_id
            )


        result = tokenizer.decode(
            output[0][inputs.input_ids.shape[-1]:],
            skip_special_tokens=True
        )


        print("\nAssistant:", result)


        if mode == "2":
            history.append({
                "role": "assistant",
                "content": result
            })


    except KeyboardInterrupt:
        print("\n\nExiting...")
        break