Update app.py

app.py

@@ -19,6 +19,9 @@ class InferencePost(BaseModel):
     question: str
     with_template: Union[str, None] = None
 
+class LLMPost(BaseModel):
+    model: str
+    question: str
 
 API_TOKEN = os.environ['HF_API_KEY']
 
@@ -50,21 +53,21 @@ Kamu adalah Asisten AI yang dikembangkan oleh Jonthan Jordan. Answer strictly in
 llama = HuggingFaceEndpoint(
     repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
     task="text-generation",
-    max_new_tokens=150,
+    max_new_tokens=4096,
     do_sample=False,
 )
 
 qwen = HuggingFaceEndpoint(
     repo_id="Qwen/Qwen1.5-4B-Chat",
     task="text-generation",
-    max_new_tokens=150,
+    max_new_tokens=4096,
     do_sample=False,
 )
 
 qwen2 = HuggingFaceEndpoint(
     repo_id="Qwen/Qwen2-1.5B-Instruct",
     task="text-generation",
-    max_new_tokens=150,
+    max_new_tokens=4096,
     do_sample=False,
 )
 
@@ -81,6 +84,19 @@ def greet_json():
     return {"Hello": "World!"}
 
 
+
+@app.post("/chat")
+async def chat(data: LLMPost):
+    if data.model == 'llama':
+        return {"data":llama.invoke(data.question)}
+    elif data.model == 'qwen':
+        return {"data":qwen.invoke(data.question)}
+    else:
+        return {"data":qwen2.invoke(data.question)}
+
+
+
+
 @app.post("/conversation")
 async def conversation(data : ConversationPost):
     return {"output":llm.invoke({"question":data.question})}