use hf api instead

app.py

@@ -6,13 +6,17 @@ from fastapi import FastAPI
 from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 
-from gradio_client import Client
+from huggingface_hub import InferenceClient
+
+client = InferenceClient(
+    "mistralai/Mistral-7B-Instruct-v0.2"
+)
 
 app = FastAPI()
 client = Client("AWeirdDev/mistral-7b-instruct-v0.2")
 
 class Message(BaseModel):
-    role: Literal["user", "assistant", "system"]
+    role: Literal["user", "assistant"]
     content: str
 
 class Payload(BaseModel):
@@ -31,6 +35,17 @@ async def stream(iter):
         except StopIteration:
             break
 
+def format_prompt(messages: List[Message]):
+    prompt = "<s>"
+    
+    for message in messages:
+        if message['role'] == 'user':
+            prompt += f"[INST] {message['content']} [/INST]"
+        else:
+            prompt += f" {message['content']}</s> "
+
+    return prompt
+
 def make_chunk_obj(i, delta, fr):
     return {
         "id": str(time.time_ns()),
@@ -49,6 +64,37 @@ def make_chunk_obj(i, delta, fr):
         ]
     }
 
+def generate(
+    messages,
+    temperature=0.9,
+    max_new_tokens=256,
+    top_p=0.95,
+    repetition_penalty=1.0,
+):
+    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=None
+    )
+
+    formatted_prompt = format_prompt(messages)
+
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+
+    for response in stream:
+        t = response.token.text
+        yield t if t != "</s>" else ""
+    
+    #return output
+
 @app.get('/')
 async def index():
     return JSONResponse({ "message": "hello", "url": "https://aweirddev-mistral-7b-instruct-v0-2-leicht.hf.space" })
@@ -68,7 +114,7 @@ async def c_cmp(payload: Payload):
                         "index": 0,
                         "message": {
                             "role": "assistant",
-                            "content": client.predict(
+                            "content": generate(
                                 payload.model_dump()['messages'],
                                 payload.temperature,
                                 4096,
@@ -85,7 +131,7 @@ async def c_cmp(payload: Payload):
         
     def streamer():
         text = ""
-        result = client.submit(
+        result = generate(
             payload.model_dump()['messages'],
             payload.temperature,	# float (numeric value between 0.0 and 1.0) in 'Temperature' Slider component
             4096,	# float (numeric value between 0 and 1048) in 'Max new tokens' Slider component