import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "olaverse/mist-qg-1.5b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
passage = "Tides are caused by the gravitational pull of the moon and, to a lesser extent, the sun, acting on Earth's oceans."
n, language = 3, "English"
messages = [
{"role": "system", "content": "You write search-style questions that a passage directly answers."},
{"role": "user", "content": f'''You are given a passage. Write {n} questions that the passage directly answers.
Rules:
- Each question must be answerable using ONLY this passage.
- Vary the type: factual, yes/no, and a comparison or "why/how".
- Natural, like a real user search query. Do NOT write "according to the passage".
- Write the questions in {language}.
Return ONLY JSON: {{"questions": ["...", "...", "..."]}}
Passage: {passage}'''},
]
input_ids = tokenizer.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
).to(model.device)
# ... see the full example on the model card