removed huggingface_hub

app.py

@@ -16,6 +16,51 @@ from llama_cpp import Llama
 import spaces
 
 
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from llama_index.core.llms import  CustomLLM, LLMMetadata, CompletionResponse, CompletionResponseGen
+from llama_index.core.llms.callbacks import llm_completion_callback
+from typing import Any
+
+
+class GemmaLLMInterface(CustomLLM):
+    model: Any
+    tokenizer: Any
+    context_window: int = 8192
+    num_output: int = 2048
+    model_name: str = "gemma_2"
+    
+    class Config:
+        protected_namespaces = () 
+
+    def _format_prompt(self, message: str) -> str:
+        return (
+            f"<start_of_turn>user\n{message}<end_of_turn>\n<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)
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        output = self.model.generate(**inputs, max_length=self.num_output)
+        raw_response = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        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:
+            yield CompletionResponse(text=token)
+
 #huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 
 

backend.py

@@ -13,6 +13,8 @@ from llama_cpp import Llama
 import spaces
 
 
+
+
 #huggingface_token = os.getenv("HUGGINGFACE_TOKEN")