Create main.py

main.py

@@ -0,0 +1,86 @@
+from flask import Flask, request, jsonify
+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")
+
+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(prompt, history, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+    # Print user prompt
+    print(f"\nUser: {prompt}\n", flush=True)
+
+    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
+    )
+
+    output = ""
+    for token in response:
+        if hasattr(token, 'token') and hasattr(token.token, 'text'):
+            output += token.token.text
+        else:
+            print(f"Unexpected token structure: {token}", flush=True)
+
+    # Print AI response
+    print(f"\nSema AI: {output}\n", flush=True)
+    return output
+
+@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:
+        response_text = generate(
+            prompt,
+            history,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty
+        )
+        return jsonify({"response": response_text})
+    except Exception as e:
+        # Print error
+        print(f"Error: {str(e)}", flush=True)
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == "__main__":
+    app.run(debug=True, port=5000)