Update app.py

app.py

@@ -4,15 +4,9 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 MODEL_ID = "devoppro/FastLLM"
-TOKENIZER_FALLBACK_ID = "Qwen/Qwen2.5-0.5B"  # FastLLM reuses the Qwen2.5 BPE vocab
-
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-except Exception as e:
-    print(f"Could not load tokenizer from {MODEL_ID} ({e}); "
-          f"falling back to {TOKENIZER_FALLBACK_ID}")
-    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_FALLBACK_ID)
+TOKENIZER_ID = "Qwen/Qwen2.5-0.5B"  # FastLLM reuses the Qwen2.5 BPE vocab
 
+tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
@@ -23,19 +17,32 @@ model.eval()
 
 
 @spaces.GPU(duration=30)
-def generate(prompt, max_new_tokens, temperature, top_p):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+def generate(prompt, max_new_tokens, temperature, top_k):
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
+    max_new_tokens = int(max_new_tokens)
+    temperature = float(temperature)
+    top_k = int(top_k)
+
     with torch.no_grad():
-        output_ids = model.generate(
-            **inputs,
-            max_new_tokens=int(max_new_tokens),
-            do_sample=True,
-            temperature=float(temperature),
-            top_p=float(top_p),
-            pad_token_id=tokenizer.eos_token_id,
-        )
-    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return full_text
+        for _ in range(max_new_tokens):
+            outputs = model(input_ids)
+            logits = outputs["logits"][:, -1, :]
+
+            logits = logits / max(temperature, 1e-5)
+
+            if top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("Inf")
+
+            probs = torch.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+
+            if next_token.item() == tokenizer.eos_token_id:
+                break
+
+    return tokenizer.decode(input_ids[0], skip_special_tokens=True)
 
 
 with gr.Blocks(title="FastLLM (150M) Demo") as demo:
@@ -43,7 +50,7 @@ with gr.Blocks(title="FastLLM (150M) Demo") as demo:
         """
         # FastLLM (150M) — Modern Causal Language Model
         A ~150M parameter decoder-only model (GQA, SwiGLU, RMSNorm, RoPE) from
-        [devoppro/FastLLM](https://huggingface.co/devoppro/FastLLM).
+        [devoppro/FastLLM](https://huggingface.co/devoppro/FastLLM), still training.
         Small model, so expect small-model quality — this is a demo, not a chatbot.
         """
     )
@@ -54,26 +61,26 @@ with gr.Blocks(title="FastLLM (150M) Demo") as demo:
                 value="Once upon a time,",
                 lines=4,
             )
-            max_new_tokens = gr.Slider(16, 512, value=128, step=16, label="Max new tokens")
+            max_new_tokens = gr.Slider(16, 256, value=60, step=8, label="Max new tokens")
             temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="Temperature")
-            top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
+            top_k = gr.Slider(0, 100, value=40, step=5, label="Top-k")
             run_btn = gr.Button("Generate", variant="primary")
         with gr.Column():
             output = gr.Textbox(label="Output", lines=12)
 
     run_btn.click(
         fn=generate,
-        inputs=[prompt, max_new_tokens, temperature, top_p],
+        inputs=[prompt, max_new_tokens, temperature, top_k],
         outputs=output,
     )
 
     gr.Examples(
         examples=[
-            ["Once upon a time,", 128, 0.7, 0.95],
-            ["The most important thing about machine learning is", 128, 0.7, 0.95],
-            ["def fibonacci(n):", 96, 0.5, 0.9],
+            ["Once upon a time,", 60, 0.7, 40],
+            ["Who are you?", 60, 0.7, 40],
+            ["The most important thing about machine learning is", 60, 0.7, 40],
         ],
-        inputs=[prompt, max_new_tokens, temperature, top_p],
+        inputs=[prompt, max_new_tokens, temperature, top_k],
     )
 
 if __name__ == "__main__":