from trulens_eval import Tru, TruCustomApp, Feedback
from trulens_eval.tru_custom_app import instrument
from trulens_eval.feedback import Coherence, Relevance
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import accelerate

class YourModel:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-1_8B-Chat", trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-1_8B-Chat", device_map="auto", trust_remote_code=True).eval()
        self.history = None

    @instrument
    def predict(self, user_input, system_prompt):
        print("Model response for:" + user_input + " with prompt:" + system_prompt)
        response, self.history = model.chat(tokenizer, user_input, history=self.history, system=system_prompt)
        return response

tru = Tru()
tru.reset_database()
model = YourModel()

coherence_feedback = Feedback(Coherence()).on_output()
relevance_feedback = Feedback(Relevance()).on_input_output()

ta = TruCustomApp(model, app_id="trumodel", feedbacks=[coherence_feedback, relevance_feedback])

def make_predictions(inputs):
    for user_input, system_prompt in inputs:
        response, record = ta.with_record(model.predict, user_input, system_prompt)
        print(response)
        print(record.dict())

inputs = [
    ("Hello!", "Hello!"),
    ("How are you?", "How are you?"),
    ("What's the weather like?", "What's the weather like?"),
    ("Tell me a joke.", "Tell me a joke."),
    ("What's new today?", "What's new today?"),
    ("Can you give me some advice?", "Can you give me some advice?"),
    ("What's your favorite color?", "What's your favorite color?"),
    ("How does AI work?", "How does AI work?"),
    ("What's the meaning of life?", "What's the meaning of life?"),
    ("Do you like music?", "Do you like music?"),
    ("What's the latest news?", "What's the latest news?"),
    ("Can you recommend a book?", "Can you recommend a book?"),
    ("What time is it?", "What time is it?"),
    ("Tell me about yourself.", "Tell me about yourself."),
    ("What's your favorite food?", "What's your favorite food?"),
    ("How do you learn?", "How do you learn?"),
    ("What are your hobbies?", "What are your hobbies?"),
    ("Can you tell a story?", "Can you tell a story?"),
    ("What is your opinion on art?", "What is your opinion on art?"),
    ("Do you know any fun facts?", "Do you know any fun facts?")
]

make_predictions(inputs)

tru.start_dashboard()