Update app.py

app.py

@@ -5,17 +5,13 @@ from huggingface_hub import InferenceClient
 import re
 import torch
 
-"""
-For more information on huggingface_hub Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-
 # Model and tokenizer loading (outside the respond function)
 try:
     tokenizer = AutoTokenizer.from_pretrained("sarvamai/sarvam-1")
     base_model = AutoModelForCausalLM.from_pretrained("sarvamai/sarvam-1")
     peft_model = PeftModel.from_pretrained(base_model, "KGSAGAR/Sarvam-1-text-normalization-3r")
     peft_model = peft_model.merge_and_unload()
-    print("Model loaded successfully!")  # Add this line
+    print("Model loaded successfully!")
 except Exception as e:
     print(f"Error loading model: {e}")
     tokenizer = None
@@ -31,7 +27,7 @@ def respond(
     top_p,
 ):
     """
-    Generates a response based on the user message and history using the provided PEFT model.
+    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).
@@ -39,13 +35,12 @@ def respond(
         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.
-    Yields:
-        str: The generated response up to the current token.
+    Returns:
+        str: The complete generated response.
     """
-    global tokenizer, peft_model #access global variables
+    global tokenizer, peft_model
     if tokenizer is None or peft_model is None:
-        yield "Model loading failed. Please check the logs."
-        return
+        return "Model loading failed. Please check the logs."
 
     # Construct the prompt
     prompt = system_message
@@ -59,19 +54,17 @@ def respond(
     # Tokenize the input prompt
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
-    # Generate the output
+    # Generate the complete output
     try:
         outputs = peft_model.generate(
             **inputs,
             max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-            do_sample=True  # Enable sampling for more diverse outputs
+            do_sample=True
         )
     except Exception as e:
-         yield f"Generation error: {e}"
-         return
-
+        return f"Generation error: {e}"
 
     # Decode the generated tokens
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -87,24 +80,17 @@ def respond(
         extracted_content = '\n'.join(match.strip() for match in matches)
         return extracted_content
 
-    # Extract the normalized text
+    # Extract and return the complete normalized text
     normalized_text = extract_user_content(generated_text)
+    return normalized_text.strip()
 
-    # Stream the response token by token
-    response = ""
-    for token in normalized_text.split():
-        response += token + " "
-        yield response.strip()
-
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
 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(
+            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.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(
@@ -117,6 +103,5 @@ demo = gr.ChatInterface(
     ],
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()