Update app.py

app.py

@@ -1,23 +1,81 @@
-from fastapi import FastAPI
-from fastapi.responses import JSONResponse
-from fastapi import Request
-from huggingface_hub import InferenceClient
+# from fastapi import FastAPI
+# from fastapi.responses import JSONResponse
+# from fastapi import Request
+# from huggingface_hub import InferenceClient
 
-app = FastAPI()
+# app = FastAPI()
 
-@app.post("/")
-async def greet_json(request: Request):
-    input_data = await request.json()
-    # number = input_data.get("number")
+# @app.post("/")
+# async def greet_json(request: Request):
+#     input_data = await request.json()
+#     # number = input_data.get("number")
  
-    # tripled_number = number * 2
-    # return {"message": f"Your input number is: {number}, your doubled number is: {tripled_number}"}
-    user_input = input_data.get("user_input")
+#     # tripled_number = number * 2
+#     # return {"message": f"Your input number is: {number}, your doubled number is: {tripled_number}"}
+#     user_input = input_data.get("user_input")
     
-    client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-    # Get the response from the model
-    response = client(user_input)
+#     client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+#     # Get the response from the model
+#     response = client(user_input)
     
-    # assistant_response = client.text_generation(user_input)
-    assistant_response = "I am assistant."
-    return {"assistant_message": f"Your input message is: {user_input}, assistant_response is: {response}"}
+#     # assistant_response = client.text_generation(user_input)
+#     assistant_response = "I am assistant."
+#     return {"assistant_message": f"Your input message is: {user_input}, assistant_response is: {response}"}
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from huggingface_hub import InferenceClient
+from typing import List, Tuple
+
+# Initialize the Inference client
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+
+# Initialize FastAPI
+app = FastAPI()
+
+class Message(BaseModel):
+    message: str
+    history: List[Tuple[str, str]]
+    system_message: str
+    max_tokens: int
+    temperature: float
+    top_p: float
+
+def generate_response(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
+    messages = [{"role": "system", "content": system_message}]
+
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+
+    for message in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = message.choices[0].delta.content
+        response += token
+    return response
+
+@app.post("/chat")
+async def chat_response(msg: Message):
+    try:
+        response = generate_response(
+            msg.message,
+            msg.history,
+            msg.system_message,
+            msg.max_tokens,
+            msg.temperature,
+            msg.top_p,
+        )
+        return {"response": response}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+