Update rag.py

rag.py

@@ -76,7 +76,7 @@ class RAG:
         print("Reranked documents")
         return documentos
         
-    def predict(self, instruction, context, model_parameters):
+    def predict_dolly(self, instruction, context, model_parameters):
 
         api_key = os.getenv("HF_TOKEN")
 
@@ -99,6 +99,34 @@ class RAG:
         response = requests.post(self.model_name, headers=headers, json=payload)
 
         return response.json()[0]["generated_text"].split("###")[-1][8:]
+
+    def predict_completion(self, instruction, context, model_parameters):
+
+        client = OpenAI(
+                base_url=ENDPOINT_URL,
+                api_key=os.getenv("HF_TOKEN")
+            )
+
+
+        chat_completion = client.chat.completions.create(
+            model="tgi",
+            messages=[
+                {"role": "user", "content": instruction}
+            ],
+            temperature=model_parameters["temperature"],
+            max_tokens=model_parameters["max_new_tokens"],
+            stream=False,
+            stop=["<|im_end|>"],
+            extra_body = {
+                "presence_penalty": model_parameters["repetition_penalty"] - 2,
+                "do_sample": False
+            }
+        )
+
+        response = chat_completion.choices[0].message.content
+
+        return response
+        
     
     def beautiful_context(self, docs):
 
@@ -122,7 +150,7 @@ class RAG:
 
             del model_parameters["NUM_CHUNKS"]
 
-            response = self.predict(prompt, text_context, model_parameters)
+            response = self.predict_completion(prompt, text_context, model_parameters)
 
             if not response:
                 return self.NO_ANSWER_MESSAGE