Update LlamaCpp import and return type

This commit is contained in:
anovazzi1 2024-01-10 10:19:38 -03:00
commit a8bf2ccf66

View file

@ -1,7 +1,7 @@
from typing import Optional, List, Dict, Any from typing import Optional, List, Dict, Any
from langflow import CustomComponent from langflow import CustomComponent
from langchain.llms import BaseLanguageModel from langchain_community.llms.llamacpp import LlamaCpp
class LlamaCppComponent(CustomComponent): class LlamaCppComponent(CustomComponent):
display_name = "LlamaCpp" display_name = "LlamaCpp"
@ -84,10 +84,10 @@ class LlamaCppComponent(CustomComponent):
use_mmap: Optional[bool] = True, use_mmap: Optional[bool] = True,
verbose: Optional[bool] = True, verbose: Optional[bool] = True,
vocab_only: Optional[bool] = False, vocab_only: Optional[bool] = False,
) -> BaseLanguageModel: ) -> LlamaCpp:
# Here you would instantiate the LlamaCpp model with the provided parameters # Here you would instantiate the LlamaCpp model with the provided parameters
# Since the actual implementation of LlamaCpp is not provided, this is a placeholder # Since the actual implementation of LlamaCpp is not provided, this is a placeholder
return BaseLanguageModel( return LlamaCpp(
model_path=model_path, model_path=model_path,
grammar=grammar, grammar=grammar,
cache=cache, cache=cache,