Update app.py

app.py

@@ -15,18 +15,29 @@ class ConversationPost(BaseModel):
     module: Union[str, None] = None
     question: str
 
+class InferencePost(BaseModel):
+    question: str
+    with_template: Union[str, None] = None
+
 
 API_TOKEN = os.environ['HF_API_KEY']
 
 os.environ["HUGGINGFACEHUB_API_TOKEN"] = API_TOKEN
 
 app = FastAPI()
-prompt = PromptTemplate.from_template("""<|im_start|>system
+prompt_qwen = PromptTemplate.from_template("""<|im_start|>system
 Kamu adalah Asisten AI yang dikembangkan oleh Jonthan Jordan. Answer strictly in Bahasa Indonesia<|im_end|>
 <|im_start|>user
 {question}<|im_end|>
 <|im_start|>assistant
 """)
+
+prompt_llama = PromptTemplate.from_template("""<|start_header_id|>system<|end_header_id|>
+
+Kamu adalah Asisten AI yang dikembangkan oleh Jonthan Jordan. Answer strictly in Bahasa Indonesia<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+""")
 # llm = prompt | HuggingFacePipeline.from_model_id(
 #     model_id="Qwen/Qwen2-1.5B-Instruct",
 #     task="text-generation",
@@ -36,20 +47,24 @@ Kamu adalah Asisten AI yang dikembangkan oleh Jonthan Jordan. Answer strictly in
 #     },
 # )
 
-llm = prompt | HuggingFaceEndpoint(
-    repo_id="Qwen/Qwen1.5-4B-Chat",
+llama = HuggingFaceEndpoint(
+    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
     task="text-generation",
     max_new_tokens=150,
     do_sample=False,
 )
 
-llm2 = prompt | HuggingFaceEndpoint(
-    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
+qwen = HuggingFaceEndpoint(
+    repo_id="Qwen/Qwen1.5-4B-Chat",
     task="text-generation",
     max_new_tokens=150,
     do_sample=False,
 )
 
+llm = prompt_qwen | qwen
+
+llm2 = prompt_llama | llama
+
 # llm = prompt | CustomLLM(repo_id="Qwen/Qwen-VL-Chat", model_type='text-generation', api_token=API_TOKEN, max_new_tokens=150).bind(stop=['<|im_end|>'])
 
 
@@ -65,4 +80,15 @@ async def conversation(data : ConversationPost):
 
 @app.post("/conversation2")
 async def conversation2(data : ConversationPost):
-    return {"output":llm2.invoke({"question":data.question})}
+    return {"output":llm2.invoke({"question":data.question})}
+
+
+@app.post("/inference")
+async def inference(data : ConversationPost):
+    if data.with_template == 'llama':
+        out = llm2.invoke(data.question)
+    elif data.with_template == 'qwen':
+        out = llm.invoke(data.question)
+    else:
+        out = llama.invoke(data.question)
+    return {"output":out}