Update app/models.py

app/models.py

@@ -3,7 +3,7 @@ import logging
 from pathlib import Path
 from typing import Optional
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoConfig, AutoModel
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
 
 import torch
 
@@ -38,7 +38,7 @@ class DataLocation(BaseModel):
                     self.cloud_uri, token=AUTH_TOKEN
                 )
 
-                model = AutoModel.from_pretrained(
+                model = AutoModelForQuestionAnswering.from_pretrained(
                     self.cloud_uri, token=AUTH_TOKEN
                 )
                 # Save the model and tokenizer locally
@@ -73,7 +73,7 @@ class QAModel:
 
         # Load the tokenizer and model from the local path
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-        self.model = AutoModel.from_pretrained(model_path,return_dict=True).to(self.device)
+        self.model = AutoModelForQuestionAnswering.from_pretrained(model_path,return_dict=True).to(self.device)
         logger.info(f"Loaded QA model: {self.model_name}")
 
     def inference_qa(self, context: str, question: str):