amusktweewt/14M-small-chat

Model Card for Model ID

Small testing version of my first model

This modelcard aims to be a base template for new models. It has been generated using this raw template.

Model Details

Model Description

Test version of my first model

Uses

Dosen't work well

Out-of-Scope Use

Better not use for anything

[More Information Needed]

Bias, Risks, and Limitations

Don't work

How to Get Started with the Model

from transformers import AutoTokenizer, AutoModelForCausalLM

# Load pre-trained model tokenizer (v3 compatibility)
tokenizer = AutoTokenizer.from_pretrained("amusktweewt/14M-small-chat")

# Load pre-trained model (PyTorch Lightning module)
model = AutoModelForCausalLM.from_pretrained("amusktweewt/14M-small-chat")

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

while True:
    user_input = input("> ")
    if user_input.lower() == "quit":
        break

    inputs = tokenizer(user_input, return_tensors="pt", max_length=512, truncation=True).to(device)
    outputs = model(**inputs)

    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=-1)
    top_prob, top_idx = torch.topk(probs, 3)  # Get the top 3 probabilities

    # Flatten the list of token IDs before decoding
    top_idx = top_idx[0].view(-1).tolist()
    top_pred = tokenizer.decode(top_idx, skip_special_tokens=True)

    print(f"You: {user_input}")
    print(f"Model: {top_pred}")

print("Goodbye!")