Spaces:
Runtime error
Runtime error
File size: 829 Bytes
dcb2a99 |
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 |
"""Simple Reasoning Agent using local models"""
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
class ReasoningAgent:
"""A simple reasoning agent using local models."""
def __init__(self, model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0"):
"""Initialize the agent with a local model."""
# Load model and tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto"
)
def get_response(self, query: str) -> str:
"""Generate a response using the local model."""
try:
# Format the prompt
prompt = f"""<|system|>
You are a helpful AI assistant.
</s>
|