Update README.md

README.md

@@ -8,7 +8,7 @@ pip install transformers[torch]
 ```
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import torch
 import time
 
@@ -19,17 +19,31 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 question = "How many Vulnerability found today"
 db_result = "243"
-summ_inp = f"Summarize the following Questions and Response: Question: {question} Response: {db_result}".lower()
+summ_inp = f"""<|im_start|>system
+Generate a clear and accurate response based on the user's question and the database output.<|im_end|>
+<|im_start|>user
+user_question:
+{question}
+db_response:
+{db_result}<|im_end|>
+<|im_start|>assistant"""
 
 import time
 start = time.time()
 
-inputs = tokenizer(summ_inp, return_tensors="pt")
+encodeds = tokenizer(summ_inp, return_tensors="pt",truncation=True).input_ids.to(device)
 model.to(device)
 inputs = inputs.to(device)
-outputs = model.generate(**inputs, max_length=526)
-answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-print(answer)
+text_streamer = TextStreamer(tokenizer, skip_prompt = True)
+response = finetuned_model.generate(
+        input_ids=encodeds,
+        streamer=text_streamer,
+        max_new_tokens=512,
+        use_cache=True,
+        pad_token_id=151645,
+        eos_token_id=151645,
+        num_return_sequences=1
+    )
 
 end = time.time()
 print(f"Time taken: {end - start}")