from pydantic import Field
from transformers import AutoModel, AutoTokenizer
from typing import Iterator
import torch

class CodegeexChatModel():
    device: str = Field(description="device to load the model")
    tokenizer = Field(description="model's tokenizer")
    model = Field(description="Codegeex model")
    temperature: float = Field(description="temperature to use for the model.")

    def __init__(self,model_name_or_path):
        super().__init__()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(
            model_name_or_path,
            trust_remote_code=True
        ).to(self.device).eval()
        print("Model has been initialized.")

    def chat(self, prompt,temperature=0.2,top_p=0.95):
        try:
            response, _ = self.model.chat(
                self.tokenizer,
                query=prompt,
                max_length=120000,
                temperature=temperature,
                top_p=top_p
            )
            return response
        except Exception as e:
            return f"error:{e}"

    def stream_chat(self,prompt,temperature=0.2,top_p=0.95):

        try:
            for response, _ in self.model.stream_chat(
                    self.tokenizer,
                    query=prompt,
                    max_length=120000,
                    temperature=temperature,
                    top_p=top_p
            ):
                yield response
        except Exception as e:
            yield f'error: {e}'