Update app/llm.py

app/llm.py

@@ -14,11 +14,11 @@ from app.users import current_active_user
 
 #from transformers import AutoModelForCausalLM
 from transformers import AutoTokenizer, pipeline
-from optimum.onnxruntime import ORTModelForQuestionAnswering
+from optimum.onnxruntime import ORTModelForCausalLM
 
-model = ORTModelForQuestionAnswering.from_pretrained("optimum/roberta-base-squad2")
-tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
 
+tokenizer = AutoTokenizer.from_pretrained("moriire/Qwen0.5-healthcare")
+model = ORTModelForCausalLM.from_pretrained("moriire/Qwen0.5-healthcare")
 class GenModel(BaseModel):
     question: str
     system: str = "You are a helpful medical AI chat assistant. Help as much as you can.Also continuously ask for possible symptoms in order to atat a conclusive ailment or sickness and possible solutions.Remember, response in English."
@@ -118,6 +118,7 @@ async def chat(chatm:ChatModel):#, user: schemas.BaseUser = fastapi.Depends(curr
 # Chat Completion API
 @llm_router.post("/generate", tags=["llm"])
 async def generate(gen:GenModel):#, user: schemas.BaseUser = fastapi.Depends(current_active_user)):
+    """
     gen.system = "You are an helpful medical AI assistant."
     gen.temperature = 0.5
     gen.seed = 42
@@ -153,5 +154,9 @@ async def generate(gen:GenModel):#, user: schemas.BaseUser = fastapi.Depends(cur
         return JSONResponse(
             status_code=500, content={"message": "Internal Server Error"}
         )
+    """
+    onnx_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    generate = onnx_gen(gen.question)
+    return generate