Update main.py

main.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
-
+from typing import List
 
 app = FastAPI()
 
@@ -10,8 +10,8 @@ client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 class Item(BaseModel):
     prompt: str
-    history: list
-    system_prompt: str
+    history: List[str] = []
+    system_prompt: str = "You are a very powerful AI assistant."
     temperature: float = 0.0
     max_new_tokens: int = 1048
     top_p: float = 0.15
@@ -26,15 +26,11 @@ def format_prompt(message, history):
     return prompt
 
 def generate(item: Item):
-    temperature = float(item.temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(item.top_p)
-
+    temperature = max(float(item.temperature), 1e-2)
     generate_kwargs = dict(
         temperature=temperature,
         max_new_tokens=item.max_new_tokens,
-        top_p=top_p,
+        top_p=float(item.top_p),
         repetition_penalty=item.repetition_penalty,
         do_sample=True,
         seed=42,
@@ -48,7 +44,23 @@ def generate(item: Item):
         output += response.token.text
     return output
 
-@app.post("/generate/")
-async def generate_text(item: Item):
+@app.get("/generate/")
+async def generate_text(
+    prompt: str,
+    history: List[str] = [],
+    system_prompt: str = "You are a very powerful AI assistant.",
+    temperature: float = 0.0,
+    max_new_tokens: int = 1048,
+    top_p: float = 0.15,
+    repetition_penalty: float = 1.0
+):
+    item = Item(
+        prompt=prompt,
+        history=history,
+        system_prompt=system_prompt,
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty
+    )
     return {"response": generate(item)}
-