Update app.py

app.py

@@ -362,22 +362,22 @@ def generate_response(user_input, model_id):
         print(f"Generated prompt: {prompt}")  # <-- Log the prompt here
         
         # Add the retrieved knowledge to the prompt
-        #func_caller.append({"role": "system", "content": prompt}) 
+        func_caller.append({"role": "system", "content": prompt}) 
 
-        #for msg in chat_history:
-            #func_caller.append({"role": "user", "content": f"{str(msg[0])}"})
-            #func_caller.append({"role": "assistant", "content": f"{str(msg[1])}"})
+        for msg in chat_history:
+            func_caller.append({"role": "user", "content": f"{str(msg[0])}"})
+            func_caller.append({"role": "assistant", "content": f"{str(msg[1])}"})
  
-        #highest_label_result = classify_intent(user_input) 
+        highest_label_result = classify_intent(user_input) 
 
         # Reformulated prompt based on intent classification
-        #reformulated_prompt = reformulate_prompt(user_input, highest_label_result)
+        reformulated_prompt = reformulate_prompt(user_input, highest_label_result)
           
-        #func_caller.append({"role": "user", "content": f'{reformulated_prompt}'})
-        #formatted_prompt = "\n".join([f"{m['role']}: {m['content']}" for m in func_caller])
+        func_caller.append({"role": "user", "content": f'{reformulated_prompt}'})
+        formatted_prompt = "\n".join([f"{m['role']}: {m['content']}" for m in func_caller])
 
         #prompt = user_input
-        #device = accelerator.device  # Automatically uses GPU or CPU based on accelerator setup
+        device = accelerator.device  # Automatically uses GPU or CPU based on accelerator setup
         
         generation_config = GenerationConfig(
             do_sample=(highest_label == "dialog continuation" or highest_label == "recommendation request"),  # True if dialog continuation, else False
@@ -394,40 +394,40 @@ def generate_response(user_input, model_id):
             )
         
         # Generate response
-        #gpt_inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
-        #gpt_output = model.generate(gpt_inputs["input_ids"], max_new_tokens=50, generation_config=generation_config)
-        #final_response = tokenizer.decode(gpt_output[0], skip_special_tokens=True)
+        gpt_inputs = tokenizer(formatted_prompt, return_tensors="pt").to(device)
+        gpt_output = model.generate(gpt_inputs["input_ids"], max_new_tokens=50, generation_config=generation_config)
+        final_response = tokenizer.decode(gpt_output[0], skip_special_tokens=True)
         # Extract AI's response only (omit the prompt)
-        #ai_response2 = final_response.replace(reformulated_prompt, "").strip()
-        #ai_response = re.sub(re.escape(formatted_prompt), "", final_response, flags=re.IGNORECASE).strip()
-        #ai_response = re.split(r'(?<=\w[.!?]) +', ai_response)
-        #ai_response = [s.strip() for s in re.split(r'(?<=\w[.!?]) +', ai_response) if s]
+        ai_response2 = final_response.replace(reformulated_prompt, "").strip()
+        ai_response = re.sub(re.escape(formatted_prompt), "", final_response, flags=re.IGNORECASE).strip()
+        ai_response = re.split(r'(?<=\w[.!?]) +', ai_response)
+        ai_response = [s.strip() for s in re.split(r'(?<=\w[.!?]) +', ai_response) if s]
        
         # Encode the prompt and candidates
-        #prompt_embedding = bertmodel.encode(formatted_prompt, convert_to_tensor=True)
-        #candidate_embeddings = bertmodel.encode(ai_response, convert_to_tensor=True)
+        prompt_embedding = bertmodel.encode(formatted_prompt, convert_to_tensor=True)
+        candidate_embeddings = bertmodel.encode(ai_response, convert_to_tensor=True)
         
         # Compute similarity scores between prompt and each candidate
-        #similarities = util.pytorch_cos_sim(prompt_embedding, candidate_embeddings)[0]
+        similarities = util.pytorch_cos_sim(prompt_embedding, candidate_embeddings)[0]
         
         # Find the candidate with the highest similarity score
         
-        #best_index = similarities.argmax()
-        #best_response = ai_response[best_index]
+        best_index = similarities.argmax()
+        best_response = ai_response[best_index]
         
         # Assuming best_response is already defined and contains the generated response
         
-        #if highest_label == "dialog continuation":
+        if highest_label == "dialog continuation":
             # Split the response into sentences
-            #sentences = best_response.split('. ')
+            sentences = best_response.split('. ')
             # Take the first three sentences and join them back together
-            #best_response = '. '.join(sentences[:3]) if len(sentences) > 3 else best_response
+            best_response = '. '.join(sentences[:3]) if len(sentences) > 3 else best_response
  
         # Append the user's message to the chat history
-        #chat_history.append({'role': 'user', 'content': user_input})
-        #chat_history.append({'role': 'assistant', 'content': best_response})
+        chat_history.append({'role': 'user', 'content': user_input})
+        chat_history.append({'role': 'assistant', 'content': best_response})
 
-        return prompt  #best_response
+        return best_response
         
     except Exception as e:
         print("Error in generate_response:")