Update app.py

app.py

@@ -1,11 +1,16 @@
 import uvicorn
-
+import os
 from flask import Flask, request, jsonify
 from huggingface_hub import InferenceClient
 
 app = Flask(__name__)
 
-API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
+# API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
+# Use environment variables for configuration
+API_URL = os.getenv("API_URL", "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1")
+# API_KEY = os.getenv("API_KEY", "your_default_api_key")  # Default API_KEY can be set here
+API_KEY = os.getenv("API_KEY") 
+
 
 def format_prompt(message, custom_instructions=None):
     prompt = ""
@@ -14,12 +19,26 @@ def format_prompt(message, custom_instructions=None):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def Mistral7B(prompt, instructions, api, temperature=0.1, max_new_tokens=2, top_p=0.95, repetition_penalty=1.0):
+def normalize_text(text):
+    # Normalize text to handle different spellings
+    replacements = {
+        'summarise': 'summarize',
+        'colour': 'color',
+        'favour': 'favor',
+        'centre': 'center',
+        # Add more replacements as needed
+    }
+    
+    text = text.lower()
+    for british, american in replacements.items():
+        text = text.replace(british, american)
+    
+    return text
+    
+def Mistral7B(prompt, instructions, api_key, temperature=0.1, max_new_tokens=2, top_p=0.95, repetition_penalty=1.0):
     global API_URL
     try:
-        temperature = float(temperature)
-        if temperature < 1e-2:
-            temperature = 1e-2
+        temperature = max(float(temperature), 1e-2)
         top_p = float(top_p)
 
         generate_kwargs = dict(
@@ -33,26 +52,27 @@ def Mistral7B(prompt, instructions, api, temperature=0.1, max_new_tokens=2, top_
         custom_instructions = instructions
         formatted_prompt = format_prompt(prompt, custom_instructions)
 
-        head = {"Authorization": f"Bearer {api}"}
-        client = InferenceClient(API_URL, headers=head)
+        client = InferenceClient(api_url=API_URL, token=api_key)
         response = client.text_generation(formatted_prompt, **generate_kwargs)
         return response
     except Exception as e:
         return str(e)
 
+
 @app.route("/generate-text", methods=["POST"])
 def generate_text():
     data = request.json
     prompt = data.get("prompt")
     instructions = data.get("instructions")
-    api_key = data.get("api_key")
+    api_key = data.get("api_key", API_KEY)  # Use provided API key or default to environment variable
 
     if not prompt or not instructions or not api_key:
         return jsonify({"error": "Missing required fields"}), 400
 
-    response = Mistral7B(prompt, instructions, api_key)
+    normalized_prompt = normalize_text(prompt)
+    response = Mistral7B(normalized_prompt, instructions, api_key)
 
     return jsonify({"response": response}), 200
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 8000)))