Update main.py

main.py

@@ -1,24 +1,13 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from huggingface_hub import InferenceClient
+
 import uvicorn
 
+from flask import Flask, request, jsonify
+from huggingface_hub import InferenceClient
 
-app = FastAPI()
+app = Flask(__name__)
 
 API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
 
-
-class Item(BaseModel):
-    prompt: str
-    history: list
-    temperature: float = 0.1
-    max_new_tokens: int = 2
-    top_p: float = 0.15
-    repetition_penalty: float = 1.0
-    instructions: str = ""
-    api: str = ""
-
 def format_prompt(message, custom_instructions=None):
     prompt = ""
     if custom_instructions:
@@ -26,32 +15,42 @@ def format_prompt(message, custom_instructions=None):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(item: Item):
-      try:
-            temperature = float(item.temperature)
-            if temperature < 1e-2:
-                  temperature = 1e-2
-            top_p = float(item.top_p)
+def Mistral7B(prompt, instructions, api, 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
+        top_p = float(top_p)
 
-            generate_kwargs = dict(
+        generate_kwargs = dict(
             temperature=temperature,
-            max_new_tokens=item.max_new_tokens,
+            max_new_tokens=max_new_tokens,
             top_p=top_p,
-            repetition_penalty=item.repetition_penalty,
+            repetition_penalty=repetition_penalty,
             do_sample=True,
-            seed=42,
-          )
-            print(item)
-            custom_instructions=item.instructions
-            formatted_prompt = format_prompt(item.prompt, custom_instructions)
-            headers = {"Authorization": f"Bearer {item.api}"}
-            client = InferenceClient(API_URL, headers=headers)
-            response = client.text_generation(formatted_prompt, **generate_kwargs)
-            return {"response": response}
-      except Exception as e:
-            return {"error": str(e)}
-
-@app.post("/generate/")
-async def generate_text(item: Item):
-    return {"response": generate(item)}
-
+            seed=69,
+        )
+        custom_instructions = instructions
+        formatted_prompt = format_prompt(prompt, custom_instructions)
+
+        head = {"Authorization": f"Bearer {api}"}
+        client = InferenceClient(API_URL, headers=head)
+        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")
+
+    if not prompt or not instructions or not api_key:
+        return jsonify({"error": "Missing required fields"}), 400
+
+    response = Mistral7B(prompt, instructions, api_key)
+
+    return jsonify({"response": response}), 200