Update funcs.py

funcs.py

@@ -93,121 +93,74 @@ def get_doc_response_emotions(user_message, therapy_session_conversation):
     print(f"User's message: {user_message}")
     print(f"RAG Matching message: {dials_embeddings.iloc[top_match_index]['Patient']}")
     print(f"Therapist's response: {dials_embeddings.iloc[top_match_index]['Doctor']}\n\n")
-
+    print(f"therapy_session_conversation: {therapy_session_conversation}")
     return '', therapy_session_conversation, emotions_msg
 
-def summarize_and_recommend(therapy_session_conversation):
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-    session_time = str(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
-    session_conversation = [item[0] for item in therapy_session_conversation]
-    session_conversation = [x for x in session_conversation if x is not None]
-    
-    session_conversation.insert(0, "Session_time: "+session_time)
-    
-    session_conversation_processed ='\n'.join(session_conversation)
-    print("session_conversation_processed:", session_conversation_processed)
+# Load model once globally for reuse
+tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
+model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct", torch_dtype=torch.float16, device_map="auto")
 
-    full_summary = ""
-    # for chunk in AI71(AI71_API_KEY).chat.completions.create(
-    #     model="tiiuae/falcon-180b-chat",
-    #     messages=[
-    #         {"role": "system", "content": """You are an Expert Cognitive Behavioural Therapist and Precis writer.
-    #         Summarize 'STRICTLY' the below user content <<<session_conversation_processed>>> 'ONLY' into useful, ethical, relevant and realistic phrases with a format
-    #       Session Time:
-    #       Summary of the patient messages: #in two to four sentences
-    #       Summary of therapist messages: #in two to three sentences:
-    #       Summary of the whole session: # in two to three sentences. Ensure the entire session summary strictly does not exceed 100 tokens."""},
-    #         {"role": "user", "content": session_conversation_processed},
-    #     ],
-    #     stream=True,
-    # ):
-    #     if chunk.choices[0].delta.content:
-    #         summary = chunk.choices[0].delta.content
-    #         full_summary += summary
-
-    from huggingface_hub import login
-    huggingface_token = os.getenv('hf_token')
-    print("hf-token:",huggingface_token)
-    login(token = huggingface_token)
-    from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
-    import torch
-    
-    # Falcon-180B model name
-    model_name = "tiiuae/falcon-7b-instruct"
-    
-    # Load tokenizer and model
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",            # Automatically maps layers to available GPUs
-        torch_dtype=torch.bfloat16,   # Use bf16 if supported
-        trust_remote_code=True,        # Falcon uses custom model code
-        # token=huggingface_token
-    )
-    
-    # Create conversation prompt (chat-style)
-    messages = [
-        {"role": "system", "content": """You are an Expert Cognitive Behavioural Therapist and Precis writer.
-    Summarize 'STRICTLY' the below user content <<<session_conversation_processed>>> 'ONLY' into useful, ethical, relevant and realistic phrases with a format
-    Session Time:
-    Summary of the patient messages: #in two to four sentences
-    Summary of therapist messages: #in two to three sentences:
-    Summary of the whole session: # in two to three sentences. Ensure the entire session summary strictly does not exceed 100 tokens."""},
-        {"role": "user", "content": session_conversation_processed}
-    ]
-    
-    # Convert to prompt string
-    prompt = ""
-    for msg in messages:
-        prompt += f"{msg['role'].capitalize()}: {msg['content']}\n"
-    prompt += "Assistant:"
-    
-    # Tokenize input
+def generate_falcon_response(prompt, max_new_tokens=300):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    # Optional: live printing during generation
-    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    
-    # Generate output
-    output = model.generate(
+    outputs = model.generate(
         **inputs,
-        max_new_tokens=200,
+        max_new_tokens=max_new_tokens,
         do_sample=True,
         temperature=0.7,
-        top_p=0.9,
-        streamer=streamer  # comment this line to just capture output
+        top_p=0.9
     )
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return decoded_output[len(prompt):].strip()
+
+def summarize_and_recommend(therapy_session_conversation):
+
+    session_time = str(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
+    session_conversation = [item[0] for item in therapy_session_conversation]
+    session_conversation = [x for x in session_conversation if x is not None]
+    session_conversation.insert(0, "Session_time: " + session_time)
     
-    # To print final output (if not using streamer)
-    print(tokenizer.decode(output[0], skip_special_tokens=True))
-    
-    
-        # full_summary = full_summary.replace('User:', '').strip()
-        # print("\n")
-        # print("Full summary:", full_summary)
-
-    full_recommendations = ""
-    for chunk in AI71(AI71_API_KEY).chat.completions.create(
-        model="tiiuae/falcon-180b-chat",
-        messages=[
-            {"role": "system", "content": """You are an expert Cognitive Behavioural Therapist.
-            Based on 'STRICTLY' the full summary <<<full_summary>>> 'ONLY' provide clinically valid, useful, appropriate action plan for the Patient as a bullted list.
-            The list shall contain both medical and non medical prescriptions, dos and donts. The format of response shall be in passive voice with proper tense.
-          - The patient is referred to........ #in one sentence
-          - The patient is advised to ........ #in one sentence
-          - The patient is refrained from........ #in one sentence
-          - It is suggested  that tha patient ........ #in one sentence
-          - Scheduled a follow-up session with the patient........#in one sentence
-            *Ensure the list contains NOT MORE THAN 7 points"""},
-            {"role": "user", "content": full_summary},
-        ],
-        stream=True,
-    ):
-        if chunk.choices[0].delta.content:
-            rec = chunk.choices[0].delta.content
-            full_recommendations += rec
-    full_recommendations = full_recommendations.replace('User:', '').strip()
-    print("\n")
-    print("Full recommendations:", full_recommendations)
-    chatbox=[]
-    return full_summary, full_recommendations, chatbox
+    session_conversation_processed = '\n'.join(session_conversation)
+    print("session_conversation_processed:", session_conversation_processed)
+
+    # Summarization prompt
+    summary_prompt = f"""You are an Expert Cognitive Behavioural Therapist and Precis writer.
+Summarize STRICTLY the following session into concise, ethical, and clinically meaningful content.
+
+Session:
+{session_conversation_processed}
+
+Format your response as:
+Session Time:
+Summary of the patient messages: 
+Summary of therapist messages: 
+Summary of the whole session: 
+Ensure the entire summary is less than 300 tokens."""
+
+    full_summary = generate_falcon_response(summary_prompt, max_new_tokens=300)
+    print("\nFull summary:", full_summary)
+
+    # Recommendation prompt
+    recommendation_prompt = f"""You are an expert Cognitive Behavioural Therapist.
+Based STRICTLY on the following summary, provide a clinically valid action plan for the patient.
+
+Summary:
+{full_summary}
+
+Use this format:
+- The patient is referred to...
+- The patient is advised to...
+- The patient is refrained from...
+- It is suggested that the patient...
+- Scheduled a follow-up session with the patient...
+
+Ensure the list contains NOT MORE THAN 7 points and is in passive voice with proper tense."""
+
+    full_recommendations = generate_falcon_response(recommendation_prompt, max_new_tokens=400)
+    print("\nFull recommendations:", full_recommendations)
+
+    chatbox = []
+    return full_summary, full_recommendations, chatbox