Spaces:
Sleeping
Sleeping
File size: 4,872 Bytes
0279c66 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import logging
import os
from typing import Optional
from groq import Groq
logger = logging.getLogger(__name__)
class ParametricMemory:
"""
Parametric Memory Module
Replaces traditional retrieval-based memory with a lightweight
parametric encoding using LoRA fine-tuning OR high-speed Groq API.
"""
def __init__(
self,
base_model_name: str = "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
lora_path: str = None,
provider: str = "local",
api_key: Optional[str] = None,
model_id: Optional[str] = None
):
self.provider = provider
self.base_model_name = base_model_name
self.lora_path = lora_path
self.device = "cuda" if torch.cuda.is_available() else "cpu"
if provider == "groq":
if not api_key:
raise ValueError("Groq API key must be provided for 'groq' provider.")
self.client = Groq(api_key=api_key)
self.model_id = model_id or "openai/gpt-oss-20b"
logger.info(f"Initialized ParametricMemory via Groq (Model: {self.model_id})")
# For compatibility with Agent loop that might check these
self.tokenizer = None
self.model = None
else:
# Local setup
logger.info(f"Loading local base model: {base_model_name} on {self.device}")
# Configure 4-bit quantization if CUDA
quantization_config = None
if torch.cuda.is_available():
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4"
)
self.tokenizer = AutoTokenizer.from_pretrained(base_model_name)
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model = AutoModelForCausalLM.from_pretrained(
base_model_name,
quantization_config=quantization_config,
device_map="auto" if self.device == "cuda" else None,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
)
if lora_path and os.path.exists(lora_path):
logger.info(f"Loading LoRA adapters from {lora_path}")
self.model = PeftModel.from_pretrained(self.model, lora_path)
def generate_solution(self, prompt: str, max_new_tokens: int = 256, temperature: float = 0.2) -> str:
"""
Generates a solution attempt using the selected provider.
"""
if self.provider == "groq":
return self._generate_groq(prompt, max_new_tokens, temperature)
else:
return self._generate_local(prompt, max_new_tokens, temperature)
def generate_reflection(self, context: str, error_signal: str, max_new_tokens: int = 256, temperature: float = 0.7) -> str:
"""
Generates reflection using the selected provider.
"""
prompt = f"""[System] You are a highly self-aware agent that reflects on its own errors.
[Context] {context}
[Error] {error_signal}
[Reflection] Analyze the error and provide a corrected reasoning path. Wrap your final answer/fix in a python code block if applicable.
"""
return self.generate_solution(prompt, max_new_tokens, temperature)
def _generate_groq(self, prompt: str, max_tokens: int, temperature: float) -> str:
try:
completion = self.client.chat.completions.create(
model=self.model_id,
messages=[{"role": "user", "content": prompt}],
temperature=temperature,
max_tokens=max_tokens,
top_p=1,
stream=False
)
return completion.choices[0].message.content.strip()
except Exception as e:
logger.error(f"Groq API Error: {e}")
return f"Error generating via Groq: {e}"
def _generate_local(self, prompt: str, max_new_tokens: int, temperature: float) -> str:
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
do_sample=True,
pad_token_id=self.tokenizer.pad_token_id
)
response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
return response.strip()
|