Update README.md

README.md

@@ -1,21 +1,89 @@
----
-language:
-- en
-pipeline_tag: text2text-generation
----
-# My Custom Model
-
-This is a custom model for text generation.
-
-## Model Details
-
-- `model_type`: Sparkoo
-
-## Usage
-
-```python
-from transformers import AutoModel, AutoTokenizer
-
-tokenizer = AutoTokenizer.from_pretrained("Sparkoo/KateAi")
-model = AutoModel.from_pretrained("Sparkoo/KateAi", from_safetensors=True)
-```
+---
+language:
+- en
+pipeline_tag: text2text-generation
+---
+# My Custom Model
+
+This is a custom model for text generation.
+
+## Model Details
+
+- `model_type`: Sparkoo
+
+## Example usage
+
+```python
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+def generate_text(prompt, model_name, max_length=100, num_return_sequences=1):
+    """
+    Generate text using the Sparkoo/KateAI model from Hugging Face Hub.
+    
+    Args:
+        prompt (str): The input text to start generation from
+        model_name (str): Name of the model on Hugging Face Hub
+        max_length (int): Maximum length of generated text
+        num_return_sequences (int): Number of different sequences to generate
+    """
+    # Load model and tokenizer
+    print(f"Loading model from {model_name}...")
+    model = GPT2LMHeadModel.from_pretrained(model_name)
+    tokenizer = GPT2Tokenizer.from_pretrained(model_name)  # Use original GPT2 tokenizer
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    # Move model to GPU if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    model.eval()
+    
+    # Encode the input prompt
+    encoded_prompt = tokenizer(prompt, return_tensors="pt", padding=True).to(device)
+    
+    # Generate text
+    print("\nGenerating text...")
+    with torch.no_grad():
+        output_sequences = model.generate(
+            input_ids=encoded_prompt["input_ids"],
+            attention_mask=encoded_prompt["attention_mask"],
+            max_length=max_length,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            do_sample=True,
+            num_return_sequences=num_return_sequences,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode and print the generated text
+    for idx, sequence in enumerate(output_sequences):
+        generated_text = tokenizer.decode(sequence, skip_special_tokens=True)
+        print(f"\nGenerated sequence {idx + 1}:")
+        print(f"{generated_text}")
+        print("-" * 50)
+
+if __name__ == "__main__":
+    # Example prompts to test
+    prompts = [
+        "Once upon a time",
+        "The artificial intelligence",
+        "In the distant future",
+        "The scientist discovered"
+    ]
+    
+    model_name = "Sparkoo/KateAI"
+    
+    # Generate text for each prompt
+    for prompt in prompts:
+        print("\n" + "="*50)
+        print(f"Prompt: {prompt}")
+        print("="*50)
+        generate_text(
+            prompt=prompt,
+            model_name=model_name,
+            max_length=200,  # Adjust as needed
+            num_return_sequences=3
+        ) 
+```