Update main.py

main.py

@@ -0,0 +1,41 @@
+from fastapi import FastAPI
+from huggingface_hub import InferenceClient
+import random
+
+API_URL = "https://api-inference.huggingface.co/models/"
+
+client = InferenceClient(
+    "mistralai/Mistral-7B-Instruct-v0.1"
+)
+
+app = FastAPI()
+
+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
+
+
+@app.post("api/v1/generate_text")
+async def generate_text(request: Request, prompt: str = Body()):
+    history = []  # You might need to handle this based on your actual usage
+
+    temperature = request.headers.get("temperature", 0.9)
+    top_p = request.headers.get("top_p", 0.95)
+    repetition_penalty = request.headers.get("repetition_penalty", 1.0)
+
+    formatted_prompt = format_prompt(prompt, history)
+    response = client.text_generation(
+        formatted_prompt,
+        temperature=temperature,
+        max_new_tokens=512,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=random.randint(0, 10**7),
+    )[0]
+
+    return response.token.text