trying keras

backend.py

@@ -14,32 +14,34 @@ from typing import Iterator, List, Any
 from llama_index.core.chat_engine import CondensePlusContextChatEngine
 from llama_index.core.llms import ChatMessage, MessageRole , CompletionResponse
 from IPython.display import Markdown, display
+import keras
+import keras_nlp
 #from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 #from llama_index import LangchainEmbedding, ServiceContext
 
+# Set the backbend before importing Keras
+os.environ["KERAS_BACKEND"] = "jax"
+# Avoid memory fragmentation on JAX backend.
+os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.00"
 
 
-huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-login(huggingface_token)
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+os.getenv("KAGGLE_USERNAME")
+os.getenv["KAGGLE_KEY"]
 
-"""model_id = "google/gemma-2-2b-it"
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto", 
-    torch_dtype= torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-    token=True)
+"""huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+login(huggingface_token)
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")"""
 
-model.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-model.eval()"""
+# Let's load Gemma using Keras
+gemma_model_id = "gemma2_instruct_2b_en"
+gemma = keras_nlp.models.GemmaCausalLM.from_preset(gemma_model_id)
 
 # what models will be used by LlamaIndex:
 Settings.embed_model = InstructorEmbedding(model_name="hkunlp/instructor-base")
 #Settings.embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'))
 
-
-Settings.llm = GemmaLLMInterface()
+#Settings.llm = GemmaLLMInterface()
+Settings.llm = GemmaLLMInterface(model=gemma)
 
 documents_paths = {
     'blockchain': 'data/blockchainprova.txt',
@@ -47,7 +49,6 @@ documents_paths = {
     'payment': 'data/paymentprova.txt'
 }
 
-
 global session_state
 session_state = {"index": False,
                  "documents_loaded": False, 
@@ -97,7 +98,6 @@ def handle_query(query_str: str,
 
     index= build_index("data/blockchainprova.txt")
     
-        
     conversation: List[ChatMessage] = []
     for user, assistant in chat_history:
       conversation.extend([
@@ -106,33 +106,6 @@ def handle_query(query_str: str,
                     ]
                 )
     
-    """if not session_state["index"]:
-        
-        matched_path = None
-        words = query_str.lower()
-        for key, path in documents_paths.items():
-            if key in words:
-                matched_path = path
-                break
-        if matched_path:
-            index = build_index(matched_path)
-            gr.Info("index costruito con la path sulla base della query")
-            session_state["index"] = True
-            
-        else: ## CHIEDI CHIARIMENTO
-          
-            conversation.append(ChatMessage(role=MessageRole.SYSTEM, content=ISTR))
-            
-            index = build_index("data/blockchainprova.txt")
-            gr.Info("index costruito con richiesta di chiarimento")
-            
-                  
-    else:
-        
-        index = build_index(matched_path)
-        #storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
-        #index = load_index_from_storage(storage_context)
-        gr.Info("index is true")"""    
 
     try:
         
@@ -175,9 +148,6 @@ def handle_query(query_str: str,
             print(info_message)
             gr.Info(info_message)"""
             
-        
-        #prompts_dict = chat_engine.get_prompts()
-        #display_prompt_dict(prompts_dict)
 
         
         #chat_engine.reset()

interface.py

@@ -6,9 +6,11 @@ import torch
 from transformers import TextIteratorStreamer
 from threading import Thread
 from pydantic import Field, field_validator
+import keras
+import keras_nlp
 
 # for transformers 2 (__setattr__ is used to bypass Pydantic check )
-class GemmaLLMInterface(CustomLLM):
+"""class GemmaLLMInterface(CustomLLM):
     def __init__(self, model_id: str = "google/gemma-2-2b-it", **kwargs):
         super().__init__(**kwargs)
         object.__setattr__(self, "model_id", model_id)  
@@ -65,5 +67,39 @@ class GemmaLLMInterface(CustomLLM):
                 yield CompletionResponse(text=streamed_response, delta=new_text)
         
         if not streamed_response:
-            yield CompletionResponse(text="No response generated.", delta="No response generated.")
+            yield CompletionResponse(text="No response generated.", delta="No response generated.")"""
+
+
+class GemmaLLMInterface(CustomLLM):
+    model: keras_nlp.models.GemmaCausalLM = None
+    context_window: int = 8192
+    num_output: int = 2048
+    model_name: str = "gemma_2"
 
+    def _format_prompt(self, message: str) -> str:
+        return (
+            f"<start_of_turn>user\n{message}<end_of_turn>\n" f"<start_of_turn>model\n"
+        )
+
+    @property
+    def metadata(self) -> LLMMetadata:
+        """Get LLM metadata."""
+        return LLMMetadata(
+            context_window=self.context_window,
+            num_output=self.num_output,
+            model_name=self.model_name,
+        )
+
+    @llm_completion_callback()
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        prompt = self._format_prompt(prompt)
+        raw_response = self.model.generate(prompt, max_length=self.num_output)
+        response = raw_response[len(prompt) :]
+        return CompletionResponse(text=response)
+
+    @llm_completion_callback()
+    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+        response = self.complete(prompt).text
+        for token in response:
+            response += token
+            yield CompletionResponse(text=response, delta=token)

requirements.txt

@@ -13,6 +13,9 @@ setuptools
 spaces
 pydantic
 ipython
+keras
+keras-nlp
+tensorflow
 #langchain
 #langchain-community
 #langchain_huggingface