Update main.py

main.py

@@ -4,6 +4,8 @@ from huggingface_hub import InferenceClient
 # Initialize Flask app
 app = Flask(__name__)
 
+print("\nHello welcome to Sema AI\n", flush=True)  # Flush to ensure immediate output
+
 # Initialize InferenceClient
 client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.1")
 
@@ -15,7 +17,9 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate_stream(prompt, history, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+def generate(prompt, history, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+    print(f"\nUser: {prompt}\n")
+
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
@@ -32,28 +36,26 @@ def generate_stream(prompt, history, temperature=0.9, max_new_tokens=256, top_p=
 
     formatted_prompt = format_prompt(prompt, history)
 
-    # Get response from Mistral model
-    response = client.text_generation(
-        formatted_prompt,
-        **generate_kwargs,
-        stream=True,
-        details=True,
-        return_full_text=False
-    )
+    try:
+        # Get response from Mistral model
+        response = client.text_generation(
+            formatted_prompt,
+            **generate_kwargs,
+            stream=True,
+            details=True,
+            return_full_text=False
+        )
 
-    def generate():
         output = ""
-        try:
-            for token in response:
-                if hasattr(token, 'token') and hasattr(token.token, 'text'):
-                    output += token.token.text
-                    yield output  # Yield intermediate response
-                else:
-                    print(f"Unexpected token structure: {token}")
-        except Exception as e:
-            print(f"Error while processing streaming response: {str(e)}")
+        for token in response:
+            output += token.token.text
+            yield token.token.text  # Yield each token for streaming
 
-    return generate
+        # Print AI response
+        print(f"\nSema AI: {output}\n")
+    except Exception as e:
+        print(f"Exception during generation: {str(e)}")
+        yield "Error occurred"
 
 @app.route("/generate", methods=["POST"])
 def generate_text():
@@ -66,7 +68,7 @@ def generate_text():
     repetition_penalty = data.get("repetition_penalty", 1.0)
 
     try:
-        return Response(stream_with_context(generate_stream(
+        return Response(stream_with_context(generate(
             prompt,
             history,
             temperature=temperature,