Update README.md

README.md

@@ -103,3 +103,84 @@ libraries, and tools for creating and deploying programs.
 
 
 ```
+
+
+# Easy Start
+
+```Python
+
+import torch
+import sys
+import warnings
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+
+warnings.filterwarnings('ignore')
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.set_float32_matmul_precision('high')
+
+MODEL_NAME = "Dront-200m"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, local_files_only=True)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    local_files_only=True,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
+model.eval()
+print("Model loaded")
+
+
+def generate_text_streaming(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+
+    # Создаем стример для потокового вывода
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    print("Output: ", end="", flush=True)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=1024,
+            temperature=0.45,
+            do_sample=True,
+            top_p=0.95,
+            repetition_penalty=1.1,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            streamer=streamer,  # Добавляем стример
+            use_cache=True
+        )
+
+    print()  # Новая строка после завершения генерации
+    return outputs
+
+
+print("Text generation started! Type 'exit' to quit.")
+while True:
+    try:
+        user_input = input("Input: ").strip()
+        if user_input.lower() in ['exit', 'quit', 'q']:
+            print("Goodbye!")
+            break
+        if not user_input:
+            continue
+
+        generate_text_streaming(user_input)
+
+    except KeyboardInterrupt:
+        print("\nGoodbye!")
+        break
+    except Exception as e:
+        print(f"Error: {e}")
+        continue
+
+```