how_to_use_the_model(inferencecode).py

# ══════════════════════════════════════════════════════════════
# FULL PIPELINE: Load Nia checkpoint + run 100-prompt eval
# ══════════════════════════════════════════════════════════════
!pip install -q transformers huggingface_hub tokenizers

import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModel
from tokenizers import Tokenizer
import csv

REPO = "Mathematicaljuice/nia"

# ── 1. Load model ──
config = AutoConfig.from_pretrained(REPO, trust_remote_code=True)
model  = AutoModel.from_config(config, trust_remote_code=True)

ckpt_path  = hf_hub_download(repo_id=REPO, filename="pytorch_model.bin")
ckpt       = torch.load(ckpt_path, map_location="cpu")
state_dict = ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt
state_dict = {k.replace("module.", "").replace("_orig_mod.", ""): v
              for k, v in state_dict.items()}
missing, unexpected = model.load_state_dict(state_dict, strict=False)
print("Missing:", missing, "| Unexpected:", unexpected)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device).eval()
print("Model class:", type(model).__name__, "| has generate():", hasattr(model, "generate"))

# ── 2. Tokenizer ──
tok_path = hf_hub_download(repo_id=REPO, filename="tokenizer.json")
tok = Tokenizer.from_file(tok_path)
IM_END = tok.token_to_id('<|im_end|>')

# ── 3. Greedy chat function (deterministic, avoids sampling noise) ──
@torch.inference_mode()
def chat_factual(instruction, max_new_tokens=80, rep_penalty=1.3):
    prompt = f'<|im_start|>user\n{instruction}\n<|im_end|>\n<|im_start|>assistant\n'
    ids = tok.encode(prompt).ids
    idx = torch.tensor([ids], dtype=torch.long, device=device)

    if hasattr(model, "generate"):
        out = model.generate(idx, max_new_tokens=max_new_tokens,
                              temperature=0.01, top_k=1, top_p=1.0,
                              repetition_penalty=rep_penalty, eos_id=IM_END)
        reply_ids = out[0].tolist()[len(ids):]
    else:
        cur = idx
        for _ in range(max_new_tokens):
            logits = model(cur)
            logits = logits[0] if isinstance(logits, tuple) else logits
            logits = logits[:, -1, :].float()
            for tid in set(cur[0].tolist()):
                logits[0, tid] /= rep_penalty
            next_tok = torch.argmax(logits, dim=-1, keepdim=True)
            cur = torch.cat([cur, next_tok], dim=1)
            if next_tok.item() == IM_END:
                break
        reply_ids = cur[0].tolist()[len(ids):]

    return tok.decode(reply_ids, skip_special_tokens=True).strip()

# ── 4. 100 test prompts ──
test_prompts = [
    # Identity (15)
    "What is your name?", "Who are you?", "Who made you?", "Who created you?",
    "Are you human?", "Are you a real person?", "What kind of AI are you?",
    "Tell me about yourself.", "What should I call you?", "Is your name Claude?",
    "Is your name GPT?", "Which company built you?", "Introduce yourself briefly.",
    "Do you have a name?", "What company are you from?",

    # General knowledge (25)
    "What is the capital of France?", "What is the capital of Malawi?",
    "What is the capital of Japan?", "Who wrote Romeo and Juliet?",
    "What is the largest planet in the solar system?", "How many continents are there?",
    "What is the boiling point of water in Celsius?", "Who was the first president of the USA?",
    "What language is spoken in Brazil?", "What is the chemical symbol for gold?",
    "How many days are in a leap year?", "What is the tallest mountain in the world?",
    "Who painted the Mona Lisa?", "What is the longest river in the world?",
    "What year did World War II end?", "What is the currency of Japan?",
    "How many bones are in the human body?", "What is the smallest country in the world?",
    "Who invented the telephone?", "What is the speed of light?",
    "What is the capital of Kenya?", "Who developed the theory of relativity?",
    "What is the largest ocean on Earth?", "What is the freezing point of water in Fahrenheit?",
    "How many strings does a standard guitar have?",

    # Plain instructions (20)
    "Give me three tips for staying focused while studying.",
    "Write a short motivational quote.", "Explain photosynthesis in one paragraph.",
    "List five healthy breakfast foods.", "Give me a recipe idea for dinner tonight.",
    "Suggest a name for a pet dog.", "Write a two-sentence bedtime story.",
    "Explain how a rainbow forms.", "Give me tips for a good night's sleep.",
    "Summarize what a computer virus is.", "Suggest three books for a beginner reader.",
    "Explain the water cycle simply.", "Give advice for a first job interview.",
    "Write a short poem about the ocean.", "Explain what gravity is.",
    "List three benefits of exercise.", "Suggest a fun weekend activity.",
    "Explain what an algorithm is.", "Give me a tip for saving money.",
    "Write a one-sentence joke.",

    # Grounded QA — answerable (15)
    "Context:\nThe Eiffel Tower was completed in 1889 and designed by engineer Gustave Eiffel.\n\nQuestion: Who designed the Eiffel Tower?\nAnswer using only the information in the context above.",
    "Context:\nMount Kilimanjaro is located in Tanzania and is the highest mountain in Africa.\n\nQuestion: Which country is Mount Kilimanjaro in?\nAnswer using only the information in the context above.",
    "Context:\nThe Great Wall of China was built over many centuries to protect Chinese states from invasions.\n\nQuestion: Why was the Great Wall of China built?\nAnswer using only the information in the context above.",
    "Context:\nMalawi is a landlocked country in southeastern Africa, bordered by Zambia, Tanzania, and Mozambique.\n\nQuestion: What type of country is Malawi?\nAnswer using only the information in the context above.",
    "Context:\nThe Amazon River flows through South America and is one of the longest rivers in the world.\n\nQuestion: Which continent does the Amazon River flow through?\nAnswer using only the information in the context above.",
    "Context:\nWilliam Shakespeare was born in Stratford-upon-Avon in 1564.\n\nQuestion: Where was William Shakespeare born?\nAnswer using only the information in the context above.",
    "Context:\nThe Sahara Desert is the largest hot desert in the world, located in North Africa.\n\nQuestion: Where is the Sahara Desert located?\nAnswer using only the information in the context above.",
    "Context:\nThe first modern Olympic Games were held in Athens, Greece, in 1896.\n\nQuestion: Where were the first modern Olympic Games held?\nAnswer using only the information in the context above.",
    "Context:\nPenicillin was discovered by Alexander Fleming in 1928.\n\nQuestion: Who discovered penicillin?\nAnswer using only the information in the context above.",
    "Context:\nThe Great Barrier Reef is located off the coast of Queensland, Australia.\n\nQuestion: Which country is the Great Barrier Reef near?\nAnswer using only the information in the context above.",
    "Context:\nThe Statue of Liberty was a gift from France to the United States in 1886.\n\nQuestion: Which country gave the Statue of Liberty to the United States?\nAnswer using only the information in the context above.",
    "Context:\nLake Malawi is one of the African Great Lakes and borders Malawi, Mozambique, and Tanzania.\n\nQuestion: What is Lake Malawi?\nAnswer using only the information in the context above.",
    "Context:\nThe printing press was invented by Johannes Gutenberg in the 15th century.\n\nQuestion: Who invented the printing press?\nAnswer using only the information in the context above.",
    "Context:\nThe Nile River flows through eleven countries in northeastern Africa.\n\nQuestion: How many countries does the Nile flow through?\nAnswer using only the information in the context above.",
    "Context:\nThe Taj Mahal was built by Mughal emperor Shah Jahan in memory of his wife.\n\nQuestion: Who built the Taj Mahal?\nAnswer using only the information in the context above.",

    # Grounded QA — unanswerable, should decline rather than guess (10)
    "Context:\nThe Eiffel Tower was completed in 1889 as the entrance arch for the World's Fair.\n\nQuestion: How tall is the Eiffel Tower?\nAnswer using only the information in the context above.",
    "Context:\nMalawi is a landlocked country in southeastern Africa.\n\nQuestion: What is the population of Malawi?\nAnswer using only the information in the context above.",
    "Context:\nWilliam Shakespeare was born in Stratford-upon-Avon.\n\nQuestion: How many plays did Shakespeare write?\nAnswer using only the information in the context above.",
    "Context:\nPenicillin was discovered by Alexander Fleming in 1928.\n\nQuestion: What disease does penicillin treat?\nAnswer using only the information in the context above.",
    "Context:\nThe first modern Olympic Games were held in Athens in 1896.\n\nQuestion: How many athletes competed in the first Olympics?\nAnswer using only the information in the context above.",
    "Context:\nThe Statue of Liberty was a gift from France in 1886.\n\nQuestion: How tall is the Statue of Liberty?\nAnswer using only the information in the context above.",
    "Context:\nThe Great Wall of China was built over many centuries.\n\nQuestion: How long is the Great Wall of China?\nAnswer using only the information in the context above.",
    "Context:\nThe printing press was invented by Johannes Gutenberg.\n\nQuestion: What city was Gutenberg born in?\nAnswer using only the information in the context above.",
    "Context:\nLake Malawi borders Malawi, Mozambique, and Tanzania.\n\nQuestion: How deep is Lake Malawi?\nAnswer using only the information in the context above.",
    "Context:\nThe Amazon River is one of the longest rivers in the world.\n\nQuestion: What animals live in the Amazon River?\nAnswer using only the information in the context above.",

    # Boundary check — identity shouldn't leak into unrelated answers (15)
    "Explain how a rainbow forms.", "What is the capital of Italy?",
    "Give me a recipe for pancakes.", "Explain what electricity is.",
    "List three planets in our solar system.", "What is 12 times 8?",
    "Explain how airplanes fly.", "Give tips for learning a new language.",
    "What is the tallest animal on Earth?", "Explain what recycling is.",
    "Suggest a good exercise routine for beginners.", "What is the capital of Germany?",
    "Explain how vaccines work.", "List three uses of solar energy.",
    "What is the largest desert in the world?",
]

print(f"\nTotal prompts: {len(test_prompts)}\n")

# ── 5. Run eval + save ──
results = []
for i, p in enumerate(test_prompts):
    reply = chat_factual(p)
    results.append({"idx": i, "prompt": p, "response": reply})
    print(f"[{i+1}/{len(test_prompts)}] {p[:60]}...")
    print(f"  -> {reply[:100]}\n")

with open("nia_eval_results.csv", "w", newline="", encoding="utf-8") as f:
    writer = csv.DictWriter(f, fieldnames=["idx", "prompt", "response"])
    writer.writeheader()
    writer.writerows(results)

print("Saved to nia_eval_results.csv")