Update huggingface_llm.py

huggingface_llm.py

@@ -9,12 +9,14 @@ class HuggingFaceLLM(LLM):
     model_id: str = Field(..., description="Hugging Face model ID")
     model: Any = Field(default=None, exclude=True)
     tokenizer: Any = Field(default=None, exclude=True)
-    device: str = Field(default="cuda" if torch.cuda.is_available() else "cpu")
     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=None, description="Device to run the model on")
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
+        if self.device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self._load_model()
 
     def _load_model(self):
@@ -48,4 +50,4 @@ class HuggingFaceLLM(LLM):
 
     @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}