Update app.py

app.py

@@ -1,17 +1,11 @@
 import gradio as gr
-import json
 import os
 from huggingface_hub import InferenceClient
-from groq import Groq
-from groq.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+from huggingface_hub.inference._generated.types.chat_completion import ChatCompletionStreamOutput
 
 # Use the fine-tuned maritime legal model
-MODEL = "nomiChroma3.1"
-##client = Groq(api_key=os.environ["GROQ_API_KEY"])
-
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-import json
+MODEL = "HuggingFaceH4/zephyr-7b-beta"
+client = InferenceClient(MODEL)
 
 def respond(
     message,
@@ -37,51 +31,30 @@ def respond(
             temperature=temperature,
             top_p=top_p,
         ):
-            print(f"Debug - Raw message type: {type(message)}")
-            print(f"Debug - Raw message content: {message}")
-            
-            if isinstance(message, str):
-                # If the message is already a string, try to parse it as JSON
-                try:
-                    parsed_message = json.loads(message)
-                    print(f"Debug - Parsed JSON: {parsed_message}")
-                    token = parsed_message.get('choices', [{}])[0].get('delta', {}).get('content', '')
-                except json.JSONDecodeError as e:
-                    print(f"JSON Decode Error: {e}")
-                    print(f"Problematic message: {message[:100]}...")  # Print first 100 chars
-                    if message.startswith('data:'):
-                        # If the message starts with 'data:', it might be SSE format
-                        clean_message = message.lstrip('data:').strip()
-                        try:
-                            parsed_message = json.loads(clean_message)
-                            print(f"Debug - Parsed JSON after cleaning: {parsed_message}")
-                            token = parsed_message.get('choices', [{}])[0].get('delta', {}).get('content', '')
-                        except json.JSONDecodeError as e:
-                            print(f"JSON Decode Error after cleaning: {e}")
-                            token = ""
-                    else:
-                        token = message  # Use the raw message as the token if parsing fails
+            if isinstance(message, ChatCompletionStreamOutput):
+                # Extract the content from the ChatCompletionStreamOutput object
+                content = message.choices[0].delta.content
+                if content is not None:
+                    response += content
+                    yield response
+                # Check if this is the last message in the stream
+                if message.choices[0].finish_reason == 'stop':
+                    break
             else:
-                # If it's not a string, it might be an object with attributes
-                try:
-                    token = message.choices[0].delta.content
-                except AttributeError:
-                    print(f"Attribute Error - message structure: {dir(message)}")
-                    token = str(message)  # Fallback to string representation
-            
-            response += token
-            yield response
+                print(f"Unexpected message type: {type(message)}")
+                print(f"Message content: {message}")
     except Exception as e:
         print(f"An error occurred: {e}")
         yield f"An error occurred: {e}"
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# Gradio interface setup
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a maritime legal assistant with expertise strictly in indian maritime law. Provide detailed legal advice and information based on indian maritime legal principles and regulations.", label="System message"),
+        gr.Textbox(
+            value="You are a maritime legal assistant with expertise strictly in Indian maritime law. Provide detailed legal advice and information based on Indian maritime legal principles and regulations.",
+            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(
@@ -93,9 +66,16 @@ demo = gr.ChatInterface(
         ),
     ],
     title="Maritime Legal Compliance",
-    description="This chatbot uses the fine tune Llama 3.1 which has the capabilities of responding and helping in legal advices regarding maritime",
+    description="This chatbot uses the Zephyr 7B Beta model to provide assistance with Indian maritime legal queries.",
+    theme="soft",
+    examples=[
+        ["What are the key regulations governing ports in India?"],
+        ["Explain the concept of cabotage in Indian maritime law."],
+        ["What are the legal requirements for registering a vessel in India?"],
+    ],
+    cache_examples=False,
 )
 
-
+# Launch the Gradio app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()