Update README.md

README.md

@@ -35,7 +35,7 @@ The Quote Generator is built on top of the GPT-2 model, fine-tuned using the Low
 
 - **Repository:** [Quote Generator](https://huggingface.co/clemsadand/quote_generator/)
 <!-- - **Paper [optional]:** N/A -->
-<!-- - **Demo [optional]:** N/A -->
+- **Demo [optional]:** N/A
 
 ## Uses
 
@@ -77,7 +77,7 @@ Use the code below to get started with the model.
 
 ```python
 from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 config = PeftConfig.from_pretrained("clemsadand/quote_generator")
 base_model = AutoModelForCausalLM.from_pretrained("gpt2")
@@ -85,10 +85,14 @@ model = PeftModel.from_pretrained(base_model, "clemsadand/quote_generator")
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
-input_text = "Generate a quote about kindness with the keywords compassion, empathy, help, generosity, care"
-input_ids = tokenizer.encode(input_text, return_tensors="pt")
+def generate_quote(input_text):
+    input_tensor = tokenizer(input_text, return_tensors="pt")
+    output = model.generate(input_tensor["input_ids"], attention_mask=input_tensor["attention_mask"],
+                          max_length=64, num_beams=5, no_repeat_ngram_size=2,
+                          early_stopping=True, pad_token_id=tokenizer.eos_token_id, do_sample=True, temperature=0.7)
+    output = tokenizer.decode(output[0], ski_special_tokens=True, clean_up_tokenization_spaces=True)
+    return output
 
-output = model.generate(input_ids)
-generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+input_text = "Generate a quote about kindness with the keywords compassion, empathy, help, generosity, care"
 
 print(generated_text)