Update README.md

README.md

@@ -1,3 +1,117 @@
 ---
 license: apache-2.0
 ---
+# NanoRush Chat
+
+NanoRush Chat is a 283M parameter GPT-style causal language model fine-tuned for conversational AI. 
+
+## Model Details & Configuration
+| Detail | Value |
+| --- | --- |
+| **Parameters** | 283M |
+| **Architecture** | GPT-2 style (Causal Language Model) |
+| **Precision** | FP16 / BFloat16 |
+| **Context Window** | Up to 4096 tokens (block_size) |
+| **Vocabulary Size** | 32,768 |
+| **Embedding Dimension (n_embd)** | 768 |
+| **Number of Heads (n_head)** | 12 |
+| **Number of Layers (n_layer)** | 36 |
+| **Base Model** | Custom pre-trained NanoRush checkpoint. |
+| **Fine-tuning Dataset** | Fine-tuned on the `HuggingFaceTB/smoltalk` dataset (a curated subset of the UltraChat 200k conversational dataset) using a supervised fine-tuning (SFT) approach. |
+| **Hardware & Optimizations** | The fine-tuning process was fully optimized for A100/H100 GPUs leveraging TF32 precision, BFloat16 autocast, and `torch.compile` for maximum throughput. |
+| **Training Strategy** | Trained using the AdamW optimizer with a cosine learning rate schedule and linear warmup. It utilizes gradient accumulation and auto-scales the batch size based on available VRAM and sequence length. |
+
+## Evaluation Results
+
+The model was evaluated using standard zero-shot accuracy metrics.
+
+| Groups / Tasks | Version | n-shot | Metric | Value | Stderr |
+|---|---|---|---|---|---|
+| **mmlu** | 2 | 0 | acc | 0.2297 | ± 0.0035 |
+| - humanities | 2 | 0 | acc | 0.2438 | ± 0.0063 |
+| - other | 2 | 0 | acc | 0.2375 | ± 0.0076 |
+| - social sciences | 2 | 0 | acc | 0.2184 | ± 0.0074 |
+| - stem | 2 | 0 | acc | 0.2119 | ± 0.0073 |
+| **arc_challenge** | 1 | 0 | acc | 0.2295 | ± 0.0123 |
+| **hellaswag** | 1 | 0 | acc | 0.3116 | ± 0.0046 |
+| **truthfulqa_mc2** | 3 | 0 | acc | 0.4320 | ± 0.0153 |
+| **winogrande** | 1 | 0 | acc | 0.5107 | ± 0.0140 |
+
+
+## Usage
+
+This model has been exported to be fully compatible with the Hugging Face `transformers` library. You can load it using the standard `AutoModelForCausalLM` pipeline.
+
+### Installation
+Make sure you have the latest version of the `transformers` and `torch` libraries installed:
+```bash
+pip install torch transformers
+```
+
+### Example Code (with Streaming & CPU Quantization)
+
+The following example demonstrates how to run NanoRush Chat efficiently on a CPU using INT8 dynamic quantization and streaming output (as used in the NanoRush web backend).
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import StoppingCriteria, StoppingCriteriaList
+from transformers.generation.streamers import TextIteratorStreamer
+import threading
+
+# Load the model and tokenizer from Hugging Face
+model_id = "Amogh1221/nano-chat"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float32, 
+    device_map="cpu"
+)
+
+# Apply INT8 dynamic quantization for CPU speedup
+print("Applying INT8 dynamic quantization...")
+model = torch.ao.quantization.quantize_dynamic(
+    model, {torch.nn.Linear}, dtype=torch.qint8
+)
+
+system_prompt = "You are NanoRush, an AI assistant created by Amogh Gupta. You are a helpful, respectful, and intelligent conversational partner. You must never pretend to be a human, and you must carefully pay attention to the conversation history."
+
+class StopOnUser(StoppingCriteria):
+    def __init__(self, prompt_length):
+        self.prompt_length = prompt_length
+    def __call__(self, input_ids, scores, **kwargs):
+        generated_tokens = input_ids[0][self.prompt_length:]
+        tail = tokenizer.decode(generated_tokens[-10:])
+        return "\nUser:" in tail or "User:" in tail
+
+# Format your prompt
+prompt = f"System: {system_prompt}\\n\\nUser: What is Quantum Computing?\\nAssistant:"
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+stop_criteria = StoppingCriteriaList([StopOnUser(prompt_length=inputs["input_ids"].shape[1])])
+streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+generation_kwargs = dict(
+    **inputs,
+    max_new_tokens=512,
+    temperature=0.7,
+    top_k=50,
+    top_p=0.9,
+    do_sample=True,
+    repetition_penalty=1.15,
+    pad_token_id=tokenizer.eos_token_id,
+    prompt_lookup_num_tokens=3,
+    stopping_criteria=stop_criteria,
+    streamer=streamer,
+)
+
+# Run generation in a background thread
+thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+thread.start()
+
+# Stream the output
+print("Assistant: ", end="")
+for text in streamer:
+    print(text, end="", flush=True)
+print()
+```