Update README.md

README.md

@@ -131,7 +131,7 @@ sequence_length = inputs["input_ids"].shape[1]
 new_output_ids = output_ids[:, sequence_length:]
 answers = tokenizer.batch_decode(new_output_ids, skip_special_tokens=True)
 
-print(answers)
+print(answers[0])
 ```
 
 
@@ -185,29 +185,18 @@ def should_filter_out(input_text, threshold=4):
 
     return max_score >= threshold
 
-def run_inference(model_path, inputs):
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    model = transformers.AutoModelForCausalLM.from_pretrained(model_path)
-    model.to(device)
-
-    tokenizer = transformers.AutoTokenizer.from_pretrained(
-        model_path,
-        model_max_length=4096,
-        padding_side="right",
-        use_fast=False,
-        add_special_tokens=False,
-    )
-    inputs = tokenizer(inputs, return_tensors='pt')
-    inputs = inputs.to(device)
-
-    output_ids = model.generate(inputs["input_ids"], max_length=4096, do_sample=False, temperature=0.0, use_cache=True)
-    sequence_length = inputs["input_ids"].shape[1]
-    new_output_ids = output_ids[:, sequence_length:]
-    answers = tokenizer.batch_decode(new_output_ids, skip_special_tokens=True)
-
-    return answers
-
 model_path = 'microsoft/Orca-2-13b'
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+model = transformers.AutoModelForCausalLM.from_pretrained(model_path)
+model.to(device)
+
+tokenizer = transformers.AutoTokenizer.from_pretrained(
+    model_path,
+    model_max_length=4096,
+    padding_side="right",
+    use_fast=False,
+    add_special_tokens=False,
+)
 
 system_message = "You are Orca, an AI language model created by Microsoft. You are a cautious assistant. You carefully follow instructions. You are helpful and harmless and you follow ethical guidelines and promote positive behavior."
 user_message = "\" \n :You can't just say, \"\"that's crap\"\" and remove it without gaining a consensus. You already know this, based on your block history. —/ \" \nIs the comment obscene? \nOptions : Yes, No."
@@ -215,7 +204,13 @@ user_message = "\" \n :You can't just say, \"\"that's crap\"\" and remove it wit
 # We use Chat Markup Language https://github.com/MicrosoftDocs/azure-docs/blob/main/articles/ai-services/openai/includes/chat-markup-language.md#working-with-chat-markup-language-chatml    
 prompt =  f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"
 
-answers = run_inference(model_path, prompt)
+inputs = tokenizer(prompt, return_tensors='pt')
+inputs = inputs.to(device)
+
+output_ids = model.generate(inputs["input_ids"], max_length=4096, do_sample=False, temperature=0.0, use_cache=True)
+sequence_length = inputs["input_ids"].shape[1]
+new_output_ids = output_ids[:, sequence_length:]
+answers = tokenizer.batch_decode(new_output_ids, skip_special_tokens=True)
 final_output = answers[0] if not should_filter_out(answers[0]) else "[Content Filtered]"
 
 print(final_output)