add stop ids

rag.py

@@ -44,7 +44,6 @@ class RAG:
 
         query = f"Context:\n{context}\n\nQuestion:\n{instruction}"
         #sys_prompt = "You are a helpful assistant. Answer the question using only the context you are provided with. If it is not possible to do it with the context, just say 'I can't answer'. <|endoftext|>"
-
         chat_completion = client.chat.completions.create(
             model="tgi",
             messages=[
@@ -53,11 +52,11 @@ class RAG:
             ],
             max_tokens=model_parameters['max_new_tokens'], # TODO: map other parameters
             frequency_penalty=model_parameters['repetition_penalty'], # this doesn't appear to do much, not a replacement for repetition penalty
-            #repetition_penalty=model_parameters['repetition_penalty'], # TODO: figure out how to pass repetition penalty
-            extra_body={'repetition_penalty': model_parameters['repetition_penalty']},
-            stream=False
+            # presence_penalty=model_parameters['repetition_penalty'],
+            # extra_body=model_parameters,
+            stream=False,
+            stop=["<|im_end|>", "<|end_header_id|>", "<|eot_id|>", "<|reserved_special_token|>"]
         )
-        print(chat_completion)
         return(chat_completion.choices[0].message.content)