Fix llamaindex ut (#10673)

* fix llamaindex ut

* add GPU ut
This commit is contained in:
Zhicun 2024-04-08 09:47:51 +08:00 committed by GitHub
parent a11b708135
commit 321bc69307
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 4 additions and 4 deletions

View file

@ -17,7 +17,7 @@
import pytest
from unittest import TestCase
import os
from ipex_llm.llamaindex.llms import BigdlLLM
from ipex_llm.llamaindex.llms import IpexLLM
class Test_LlamaIndex_Transformers_API(TestCase):
def setUp(self):
@ -50,7 +50,7 @@ class Test_LlamaIndex_Transformers_API(TestCase):
return prompt
def test_bigdl_llm(self):
llm = BigdlLLM(
llm = IpexLLM.from_model_id(
model_name=self.llama_model_path,
tokenizer_name=self.llama_model_path,
context_window=512,

View file

@ -18,7 +18,7 @@ import torch
import pytest
from unittest import TestCase
import os
from ipex_llm.llamaindex.llms import BigdlLLM
from ipex_llm.llamaindex.llms import IpexLLM
class Test_LlamaIndex_Transformers_API(TestCase):
def setUp(self):
@ -51,7 +51,7 @@ class Test_LlamaIndex_Transformers_API(TestCase):
return prompt
def test_bigdl_llm(self):
llm = BigdlLLM(
llm = IpexLLM.from_model_id(
model_name=self.llama_model_path,
tokenizer_name=self.llama_model_path,
context_window=512,