Update app.py

app.py

@@ -133,8 +133,7 @@ class MistralChatBot:
     def __init__(self, system_prompt="Le dialogue suivant est une conversation"):
         self.system_prompt = system_prompt
 
-    def predict(self, user_message, context):
-        #detailed_prompt = """### Question ###\n""" + user_message + "\n\n### Contexte ###\n" + context + "\n\n### Formule ###\n"
+    def predict_field(self, user_message, context):
         detailed_prompt = """### Question ###\n""" + user_message + "\n\n### Field ###\n"
         prompts = [detailed_prompt]
         outputs = llm.generate(prompts, sampling_params, use_tqdm = False)
@@ -143,6 +142,15 @@ class MistralChatBot:
         fiches_html = ""
         return generated_text, fiches_html
 
+    def predict_answer(self, user_message, context):
+        detailed_prompt = """### Question ###\n""" + user_message + "\n\n### Contexte ###\n" + context + "\n\n### Formule ###\n"
+        prompts = [detailed_prompt]
+        outputs = llm.generate(prompts, sampling_params, use_tqdm = False)
+        generated_text = outputs[0].outputs[0].text
+        print(generated_text)
+        fiches_html = ""
+        return generated_text, fiches_html
+
 # Create the Falcon chatbot instance
 mistral_bot = MistralChatBot()