Update app.py

app.py

@@ -1,21 +1,15 @@
-from flask import Flask, render_template, request, jsonify
+from flask import Flask, request, jsonify
 from huggingface_hub import InferenceClient
-import gradio as gr
 
 app = Flask(__name__)
+app.config["DEBUG"] = True  # Enable for debugging
 
+# Load model client
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
-def format_prompt(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]"
-    return prompt
-
+# Function for text generation with enhanced prompt formatting
 def generate(
-    prompt, history, system_prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,
+    prompt, history, system_prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0
 ):
     temperature = float(temperature)
     if temperature < 1e-2:
@@ -31,29 +25,34 @@ def generate(
         seed=42,
     )
 
-    formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+    # Enhanced prompt formatting for better context
+    formatted_prompt = f"{system_prompt}\n{''.join(f'{user_prompt} ||| {bot_response}\n' for user_prompt, bot_response in history)}\n{prompt}"
+
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
 
     for response in stream:
         output += response.token.text
-        yield output
     return output
 
-@app.route('/generate', methods=['POST'])
-def generate_response():
-    data = request.get_json()
-    prompt = data['prompt']
-    history = data['history']
-    system_prompt = data['system_prompt']
-    temperature = data['temperature']
-    max_new_tokens = data['max_new_tokens']
-    top_p = data['top_p']
-    repetition_penalty = data['repetition_penalty']
 
-    result = list(generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty))
+@app.route("/generate", methods=["POST"])
+def generate_text():
+    data = request.json
+    prompt = data.get("prompt")
+    history = data.get("history", [])
+    system_prompt = data.get("system_prompt")
+    temperature = data.get("temperature", 0.9)
+    max_new_tokens = data.get("max_new_tokens", 256)
+    top_p = data.get("top_p", 0.95)
+    repetition_penalty = data.get("repetition_penalty", 1.0)
+
+    response = generate(
+        prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty
+    )
+
+    return jsonify({"response": response})
 
-    return jsonify({'result': result})
 
-if __name__ == '__main__':
-    app.run(port=7860)
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)