llm folder added

llm/__init__.py

@@ -0,0 +1 @@
+from .model_registry import get_llm

llm/llm_base.py

@@ -0,0 +1,3 @@
+class LLMBase:
+    def generate(self, prompt: str, max_new_tokens: int = 200, **kwargs) -> str:
+        raise NotImplementedError("Must implement in subclass.")

llm/model_registry.py

@@ -0,0 +1,13 @@
+from llm.llm_base import LLMBase
+from transformers import pipeline
+
+class LocalLLM(LLMBase):
+    def __init__(self, model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device="cpu", **kwargs):
+        self.pipe = pipeline("text-generation", model=model_name, device=device, **kwargs)
+    def generate(self, prompt: str, max_new_tokens: int = 200, **kwargs) -> str:
+        result = self.pipe(prompt, max_new_tokens=max_new_tokens)
+        return result[0]["generated_text"]
+
+def get_llm(model_name: str = "TinyLlama/TinyLlama-1.1B-Chat-v1.0", **kwargs):
+    
+    return LocalLLM(model_name=model_name, **kwargs)