Update app.py

app.py

@@ -41,7 +41,24 @@ def format_prompt(message, history):
   prompt += f"[INST] {message} [/INST]"
   return prompt
 
-def generate_response(prompt, history, temperature=0.9, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0):
+pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200, do_sample=True,
+    max_new_tokens=1024, 
+    temperature=0.9, 
+    top_k=50, 
+    top_p=0.95,
+    num_return_sequences=1)
+
+def generate_response(message, history):
+    prompt = "<s>"
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    result = pipe(f"{prompt}")[0]['generated_text']
+    return result
+
+'''
+def generate_response(prompt, history, temperature=0.9, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0, **kwargs,):
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
@@ -58,6 +75,22 @@ def generate_response(prompt, history, temperature=0.9, max_new_tokens=1024, top
     runtimeFlag = "cuda:0"
     formatted_prompt = format_prompt(f"{prompt}", history)
     inputs = tokenizer([formatted_prompt], return_tensors="pt").to(runtimeFlag)
+    generation_config = GenerationConfig(
+        temperature=temperature,
+        top_p=top_p,
+        max_new_tokens=max_new_tokens,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        **kwargs,
+    )
+    generation_output = model.generate(
+            **inputs,
+            generation_config=generation_config,
+            return_dict_in_generate=True,
+            output_scores=True,
+            max_new_tokens=max_new_tokens,
+        )
+'''
 
 # UI design
 examples=[
@@ -67,7 +100,7 @@ examples=[
 ]
 
 gr.ChatInterface(
-    fn=generate,
+    fn=generate_response,
     chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
     title="Psychological Assistant: Expert in Assessment and Strategic Planning",
     description="Enter counseling notes to generate an assessment and plan.",