Update app/llm.py

app/llm.py

@@ -12,11 +12,12 @@ from fastapi import APIRouter
 from app.users import current_active_user
 
 
-from transformers import AutoModelForCausalLM
-
-model = AutoModelForCausalLM.from_pretrained("bigcode/starcoder")
-model.to_bettertransformer()
+#from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer, pipeline
+from optimum.onnxruntime import ORTModelForQuestionAnswering
 
+model = ORTModelForQuestionAnswering.from_pretrained("optimum/roberta-base-squad2")
+tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
 
 class GenModel(BaseModel):
     question: str
@@ -35,6 +36,7 @@ class ChatModel(BaseModel):
     mirostat_mode: int=2
     mirostat_tau: float=4.0
     mirostat_eta: float=1.1
+    
 llm_chat = llama_cpp.Llama.from_pretrained(
     repo_id="Qwen/Qwen1.5-0.5B-Chat-GGUF",
     filename="*q4_0.gguf",
@@ -57,6 +59,7 @@ llm_generate = llama_cpp.Llama.from_pretrained(
     mirostat_eta=1.1,
     #chat_format="llama-2"
 )
+
 # Logger setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -82,7 +85,11 @@ def health():
 # Chat Completion API
 @llm_router.post("/chat/", tags=["llm"])
 async def chat(chatm:ChatModel):#, user: schemas.BaseUser = fastapi.Depends(current_active_user)):
-    
+    onnx_qa = pipeline("question-answering", model=model, tokenizer=tokenizer)
+    question = "What's my name?"
+    context = "My name is Philipp and I live in Nuremberg."
+    pred = onnx_qa(question, context)
+    return pred
     """
     #chatm.system = chatm.system.format("")#user.email)
     try: