Update README.md

README.md

@@ -15,6 +15,7 @@ pipeline_tag: text-generation
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import time
 
 MODEL_NAME = "VortexIntelligence/VLM-1.1-K1-Preview"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -23,26 +24,52 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
 tokenizer.add_special_tokens({'additional_special_tokens': ['<|system|>', '<|user|>', '<|assistant|>']})
 tokenizer.eos_token = "<|endoftext|>"
 model.resize_token_embeddings(len(tokenizer))
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
-
-def generate_response(user_input):
+def stream_response(user_input):
     system_prompt = "You are a helpful assistant."
     prompt = f"<|system|>\n{system_prompt}\n<|user|>\n{user_input}\n<|assistant|>\n"
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True).to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    input_len = inputs['input_ids'].shape[-1]
+    max_new_tokens = 128
 
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=128,
-        do_sample=False,
-        eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.pad_token_id,
-        repetition_penalty=2.1
-    )
+    start_time = time.time()
+
+    streamer = []
+    for i in range(1, max_new_tokens + 1):
+        output = model.generate(
+            **inputs,
+            max_new_tokens=i,
+            do_sample=False,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+            repetition_penalty=2.1
+        )
+        output_tokens = output[0][input_len:]
+        generated_text = tokenizer.decode(output_tokens, skip_special_tokens=True)
+        
+        if len(streamer) < len(generated_text):
+            new_chunk = generated_text[len(streamer):]
+            print(new_chunk, end='', flush=True)
+            streamer += new_chunk
+        if tokenizer.eos_token in tokenizer.decode(output[0]):
+            break
 
-    generated_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
-    return generated_text.strip()
+    end_time = time.time()
+    duration = end_time - start_time
+    total_tokens = len(streamer)
+    tps = total_tokens / duration
+    tpm = tps * 60
+
+    print("\n")
+    print("-"*20)
+    print(f"Time taken: {duration:.2f}s")
+    print(f"Total tokens: {total_tokens}")
+    print(f"Tokens/sec: {tps:.2f}")
+    print(f"Tokens/min: {tpm:.2f}")
 
 print("VLM 1.1 Chat - Type 'exit' to quit")
 while True:
@@ -50,9 +77,6 @@ while True:
     if user_input.lower() == 'exit':
         print("Exiting chat. Goodbye!")
         break
-        
-    assistant_response = generate_response(user_input)
-
-    print(f"VLM: {assistant_response}")
-
+    print("VLM: ", end="", flush=True)
+    stream_response(user_input)
 ```