Update huggingface_llm.py

huggingface_llm.py

@@ -9,18 +9,20 @@ class HuggingFaceLLM(LLM):
     model_id: str = Field(..., description="Hugging Face model ID")
     temperature: float = Field(default=0.7, description="Sampling temperature")
     max_tokens: int = Field(default=256, description="Maximum number of tokens to generate")
+    device: str = Field(default="cpu", description="Device to run the model on")
+    
     _model: Any = PrivateAttr()
     _tokenizer: Any = PrivateAttr()
-    # _device: str = PrivateAttr()
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        # self._device = "cpu"
+        if self.device == "cpu":
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self._load_model()
 
     def _load_model(self):
         self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
-        self._model = AutoModelForCausalLM.from_pretrained(self.model_id)
+        self._model = AutoModelForCausalLM.from_pretrained(self.model_id).to(self.device)
 
     @property
     def _llm_type(self) -> str:
@@ -33,8 +35,8 @@ class HuggingFaceLLM(LLM):
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> str:
-        input_ids = self._tokenizer.encode(prompt, return_tensors="pt")
-        
+        input_ids = self._tokenizer.encode(prompt, return_tensors="pt").to(self.device)
+
         with torch.no_grad():
             output = self._model.generate(
                 input_ids,
@@ -43,10 +45,10 @@ class HuggingFaceLLM(LLM):
                 do_sample=True,
                 pad_token_id=self._tokenizer.eos_token_id
             )
-        
+
         response = self._tokenizer.decode(output[0], skip_special_tokens=True)
         return response[len(prompt):].strip()
 
     @property
     def _identifying_params(self) -> Dict[str, Any]:
-        return {"model_id": self.model_id, "temperature": self.temperature, "max_tokens": self.max_tokens}
+        return {"model_id": self.model_id, "temperature": self.temperature, "max_tokens": self.max_tokens, "device": self.device}