Update README.md

README.md

@@ -50,5 +50,129 @@ You can find the full code in this repo as `train.py` and inference.py. Have fun
 ## Usage
 Use this to run the model:
 ```python3
+"""
+StorySupra-10M — Interactive Story Generator
+Loads model weights directly from HuggingFace: SupraLabs/StorySupra-10M
+"""
 
+import torch
+from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
+
+# ──────────────────────────────────────────────
+# Configuration
+# ──────────────────────────────────────────────
+MODEL_ID = "SupraLabs/StorySupra-10M"
+
+GENERATION_DEFAULTS = {
+    "max_new_tokens": 100,
+    "temperature": 0.55,
+    "top_k": 25,
+    "top_p": 0.85,
+    "repetition_penalty": 1.1,
+    "do_sample": True,
+}
+
+EXIT_COMMANDS = {"exit", "quit", "leave"}
+
+# ──────────────────────────────────────────────
+# Model loading
+# ──────────────────────────────────────────────
+
+def load_model(model_id: str):
+    """Download and return the tokenizer and model from HuggingFace Hub."""
+    print(f"Downloading model from HuggingFace: {model_id}")
+    print("(This may take a moment on first run — weights will be cached locally.)\n")
+
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id)
+    model = LlamaForCausalLM.from_pretrained(model_id)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}\n")
+
+    model.to(device)
+    model.eval()
+
+    return tokenizer, model, device
+
+
+# ──────────────────────────────────────────────
+# Text generation
+# ──────────────────────────────────────────────
+
+def generate_text(
+    prompt: str,
+    tokenizer,
+    model,
+    device: str,
+    max_new_tokens: int = GENERATION_DEFAULTS["max_new_tokens"],
+    temperature: float = GENERATION_DEFAULTS["temperature"],
+    top_k: int = GENERATION_DEFAULTS["top_k"],
+    top_p: float = GENERATION_DEFAULTS["top_p"],
+    repetition_penalty: float = GENERATION_DEFAULTS["repetition_penalty"],
+) -> str:
+    """Generate a story continuation from the given prompt."""
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        output_tokens = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+
+
+# ──────────────────────────────────────────────
+# Interactive loop
+# ──────────────────────────────────────────────
+
+def run():
+    print("=" * 50)
+    print("  StorySupra-10M — Interactive Story Generator")
+    print("=" * 50)
+
+    tokenizer, model, device = load_model(MODEL_ID)
+
+    print("-" * 50)
+    print("Model ready! Type a prompt to generate a story.")
+    print(f"Type {' / '.join(EXIT_COMMANDS)} to quit.")
+    print("-" * 50)
+
+    while True:
+        try:
+            user_prompt = input("\nYour prompt: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nExiting. Goodbye!")
+            break
+
+        if not user_prompt:
+            print("Please enter a prompt.")
+            continue
+
+        if user_prompt.lower() in EXIT_COMMANDS:
+            print("Goodbye!")
+            break
+
+        print("\nGenerating...\n")
+        story = generate_text(user_prompt, tokenizer, model, device)
+
+        print("Generated story:")
+        print("-" * 20)
+        print(story)
+        print("-" * 20)
+
+
+# ──────────────────────────────────────────────
+# Entry point
+# ──────────────────────────────────────────────
+
+if __name__ == "__main__":
+    run()
 ```