Update app.py

app.py

@@ -18,43 +18,24 @@ except Exception as e:
     base_model = None
     peft_model = None
 
-def respond(
-    message,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    """
-    Generates a complete response based on the user message and history using the provided PEFT model.
-    Args:
-        message (str): The user's input message.
-        history (list of tuples): A list containing tuples of (user_message, assistant_response).
-        system_message (str): The system's initial message or prompt.
-        max_tokens (int): The maximum number of tokens to generate.
-        temperature (float): The temperature parameter for generation.
-        top_p (float): The top_p parameter for nucleus sampling.
-    Returns:
-        str: The complete generated response.
-    """
+def respond(message, history, system_message, max_tokens, temperature, top_p):
     global tokenizer, peft_model
+    
     if tokenizer is None or peft_model is None:
         return "Model loading failed. Please check the logs."
-
+    
     # Construct the prompt
     prompt = system_message
     for user_msg, assistant_msg in history:
         if user_msg:
-            prompt += f"<user>{user_msg}</user>"
+            prompt += f"\nUser: {user_msg}"
         if assistant_msg:
-            prompt += f"<assistant>{assistant_msg}</assistant>"
-    prompt += f"<user>{message}</user>"
-
+            prompt += f"\nAssistant: {assistant_msg}"
+    prompt += f"\nUser: {message}"
+    
     # Tokenize the input prompt
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-
-    # Generate the complete output
+    
     try:
         outputs = peft_model.generate(
             **inputs,
@@ -63,43 +44,24 @@ def respond(
             top_p=top_p,
             do_sample=True
         )
+        
+        # Get the last message from the generated text
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        lines = generated_text.split('\n')
+        last_line = lines[-1] if lines else ""
+        
+        return last_line.strip()
+        
     except Exception as e:
         return f"Generation error: {e}"
 
-    # Decode the generated tokens
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Extract content between <user>...</user> tags
-    def extract_user_content(text):
-        """
-        Extracts and returns content between <user>...</user> tags in the given text.
-        If multiple such sections exist, their contents are concatenated.
-        """
-        pattern = r'<user>(.*?)</user>'
-        matches = re.findall(pattern, text, re.DOTALL)
-        extracted_content = '\n'.join(match.strip() for match in matches)
-        return extracted_content
-
-    # Extract and return the complete normalized text
-    normalized_text = extract_user_content(generated_text)
-    return normalized_text.strip()
-
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(
-            value="Take the user input in Hindi language and normalize specific entities, Only including: Dates (any format) Currencies Scientific units, <Example> Exampleinput :  2012–13 में रक्षा सेवाओं के लिए 1,93,407 करोड़ रुपए का प्रावधान किया गया था, जबकि 2011–2012 में यह राशि 1,64,415 करोइ़ थी, Exampleoutput: ट्वेन्टी ट्वेल्व फिफ्टीन में रक्षा सेवाओं के लिए वन करोड़ निनेटी थ्री थाउजेंड फोर हंड्रेड सेवन करोड़ रुपए का प्रावधान किया गया था, जबकि ट्वेन्टी एलेवन ट्वेल्व में यह राशि वन करोड़ सिक्स्टी फोर थाउजेंड फोर हंड्रेड फिफ्टीन करोड़ थी </Example>, Only provide the normalized output with atmost accuracy <user> input:",
-            label="System message"
-        ),
+        gr.Textbox(label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
     ],
 )