from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

class DeepSeek:
    def __init__(self):
        model_id = "deepseek-ai/deepseek-llm-7b-chat"
        offload_dir = "./offload/deepseek-chat"
        os.makedirs(offload_dir, exist_ok=True)

        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="auto",
            offload_folder=offload_dir
        ).eval()

    def generate(self, prompt, temperature=0.7):
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        outputs = self.model.generate(
            **inputs,
            do_sample=True,
            temperature=temperature,
            top_p=0.9,
            max_new_tokens=128,
            no_repeat_ngram_size=2,
            eos_token_id=self.tokenizer.eos_token_id
        )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)