Spaces:
Runtime error
Runtime error
File size: 1,050 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 27 |
"""Local reasoning agent using TinyLlama."""
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
class ReasoningAgent:
"""Local reasoning agent using TinyLlama with chain of thought prompting."""
def __init__(self, model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0"):
"""Initialize the agent with 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, message: str) -> str:
"""Generate response using local model with chain of thought prompting."""
try:
# Format prompt with chain of thought structure
prompt = f"""<|system|>
You are a helpful AI assistant that uses chain of thought reasoning to answer questions.
For each response, break down your thinking into steps before giving the final answer.
</s>
|