Update README.md

README.md

@@ -15,14 +15,15 @@ pipeline_tag: text-generation
 trust_remote_code: true
 special_tokens:
   additional_special_tokens:
-    - "[شروع_شعر]"
-    - "[پایان_شعر]"
-    - "[مصرع]"
+  - '[شروع_شعر]'
+  - '[پایان_شعر]'
+  - '[مصرع]'
 quantization:
   load_in_4bit: true
   bnb_4bit_quant_type: nf4
   bnb_4bit_compute_dtype: float16
   bnb_4bit_use_double_quant: true
+license: mit
 ---
 
 # Model Card for llama_poetry_fa
@@ -118,16 +119,122 @@ Users should critically evaluate the generated text, especially if used in publi
 
 ## How to Get Started with the Model
 
-Use the provided `poetry_generator.py` script:
+## Poetry Generator Code
 
-```python
-from poetry_generator import PoetryGenerator
+```bash
+pip install -U transformers>=4.30.0
+pip install -U accelerate
+pip install bitsandbytes==0.42.0
+```
 
-generator = PoetryGenerator(
-    model_path="/path/to/your/trained/model",
-    token="your_hf_token"
-)
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import PeftModel
+
+class PoetryGenerator:
+    def __init__(self, model_path, token):
+        self.token = token
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Configure quantization settings
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True
+        )
+
+        # Load tokenizer from the base model used during fine-tuning
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            "meta-llama/Llama-3.1-8B-Instruct",
+            token=token,
+            trust_remote_code=True
+        )
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        # Add the special tokens that were used during training
+        special_tokens = {
+            "additional_special_tokens": [
+                "[شروع_شعر]",
+                "[پایان_شعر]",
+                "[مصرع]"
+            ]
+        }
+        self.tokenizer.add_special_tokens(special_tokens)
+
+        # Load the base model
+        base_model = AutoModelForCausalLM.from_pretrained(
+            "meta-llama/Llama-3.1-8B-Instruct",
+            token=token,
+            device_map="auto",
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+            quantization_config=bnb_config
+        )
+
+        # Resize token embeddings to match tokenizer
+        base_model.resize_token_embeddings(len(self.tokenizer))
+
+        # Load the fine-tuned model from Hugging Face Hub
+        self.model = PeftModel.from_pretrained(
+            base_model,
+            model_path,
+            token=token,
+            device_map="auto"
+        )
+
+        self.model.eval()
+
+    def generate_poem(self, prompt):
+        formatted_prompt = f"""سوال: {prompt}
+لطفا یک شعر فارسی در پاسخ به این سوال بسرایید که دارای وزن و قافیه مناسب باشد.
+
+شعر:"""
+
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt", padding=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_length=512,
+                num_return_sequences=1,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id
+            )
+
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def main():
+    # Use the Hugging Face Hub model path instead of a local path
+    generator = PoetryGenerator(
+        model_path="8lianno/llama_poetry_fa",
+        token="<YOUR_HF_TOKEN>"
+    )
+
+    prompts = [
+        "درباره بهار شعری بسرایید",
+        "شعری درباره عشق بنویسید",
+        "درباره دریا شعری بسرایید"
+    ]
+
+    print("=== Persian Poetry Generation ===\n")
+    for i, prompt in enumerate(prompts, 1):
+        print(f"\nPrompt {i}: {prompt}")
+        print("\nGenerated Poetry:")
+        try:
+            poem = generator.generate_poem(prompt)
+            print(poem)
+            print("\n" + "="*50)
+        except Exception as e:
+            print(f"Error generating poem: {str(e)}")
+            print(f"Error type: {type(e)}")
+
+if __name__ == "__main__":
+    main()
+```
 
-prompt = "امید چیست؟"
-poem = generator.generate_poem(prompt)
-print(poem)