Update src/prediction.py

src/prediction.py

@@ -1,4 +1,4 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, AutoConfig
 from textwrap import dedent
 from build_db_index import query_vectorstore
 from langdetect import detect
@@ -88,7 +88,9 @@ class Prediction:
         return prediction
     
     def configure_model(self, model_name):
-        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+        config = AutoConfig.from_pretrained(model_name)
+        config.model_type = "qwen2"
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
         self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
         # self.tokenizer.chat_template = """{% for message in messages %}
         #         {% if message['role'] == 'system' -%}