Create main.py

main.py

@@ -0,0 +1,82 @@
+from flask import Flask, request, jsonify, Response, stream_with_context
+from huggingface_hub import InferenceClient
+
+# Initialize Flask app
+app = Flask(__name__)
+
+# Initialize InferenceClient
+client = InferenceClient("mistralai/Mistral-7B-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
+
+def generate_stream(prompt, history, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42,
+    )
+
+    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
+    )
+
+    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)}")
+
+    return generate
+
+@app.route("/generate", methods=["POST"])
+def generate_text():
+    data = request.json
+    prompt = data.get("prompt", "")
+    history = data.get("history", [])
+    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)
+
+    try:
+        return Response(stream_with_context(generate_stream(
+            prompt,
+            history,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty
+        )), content_type='text/plain')
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == "__main__":
+    app.run(debug=True, port=5000)