src/backend/model_operations.py

import os
import time
from datetime import datetime
import logging
from pathlib import Path  
import requests
import json

import numpy as np
import pandas as pd
import spacy
import litellm
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, AutoModelForTokenClassification, AutoConfig, Qwen2VLForConditionalGeneration, AutoProcessor
from peft import PeftModel
import torch
import cohere
from openai import OpenAI
from together import Together
import anthropic
import replicate
# import google.generativeai as genai
import vertexai
from vertexai.generative_models import GenerativeModel, Part, SafetySetting, FinishReason
from mistralai import Mistral
from qwen_vl_utils import process_vision_info


import src.backend.util as util
import src.envs as envs

litellm.set_verbose=True

# Set up basic configuration for logging
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')

# Load spacy model for word tokenization
nlp = spacy.load("en_core_web_sm")

os.environ["HUGGINGFACE_API_KEY"] =  envs.TOKEN

class ModelLoadingException(Exception):
    """Exception raised for errors in loading a model.

    Attributes:
        model_id (str): The model identifier.
        revision (str): The model revision.
    """

    def __init__(self, model_id, revision, messages="Error initializing model"):
        self.model_id = model_id
        self.revision = revision
        super().__init__(f"{messages} id={model_id} revision={revision}")


class SummaryGenerator:
    """A class to generate summaries using a causal language model.

    Attributes:
        model (str): huggingface/{model_id}
        api_base (str): https://api-inference.huggingface.co/models/{model_id}
        summaries_df (DataFrame): DataFrame to store generated summaries.
        revision (str): Model revision.
        avg_length (float): Average length of summaries.
        answer_rate (float): Rate of non-empty summaries.
    """

    def __init__(self, model_id, revision, device):
        """
        Initializes the SummaryGenerator with a model.

        Args:
            model_id (str): Identifier for the model.
            revision (str): Revision of the model.
        """
        self.model_id = model_id
        self.model = f"huggingface/{model_id}"
        self.api_base = f"https://api-inference.huggingface.co/models/{model_id}"
        self.summaries_df = pd.DataFrame()
        self.revision = revision
        self.device = device
        self.avg_length = None
        self.answer_rate = None
        self.exceptions = None
        self.local_model = None
        self.local_pipeline = None

    def generate_summaries(self, df, save_path=None):
        """Generate summaries for a given DataFrame of source docs.

        Args:
            df (DataFrame): DataFrame containing source docs.

        Returns:
            summaries_df (DataFrame): Generated summaries by the model.
        """
        exceptions = []
        if (save_path is not None) and os.path.exists(save_path):
            self.summaries_df = pd.read_csv(save_path)
            print(f'Loaded generated summaries from {save_path}')
        else:
            source, summary, dataset = [], [], [] 
            print(f"Total: {df.shape[0]}")
            for index, row in tqdm(df.iterrows(), total=df.shape[0]):
                _source = row['text']
                _dataset = row['dataset']

                system_prompt = envs.SYSTEM_PROMPT
                user_prompt = f"{envs.USER_PROMPT}\nPassage:\n{_source}"
                _summary = None

                while not _summary:
                    try:
                        _summary = self.generate_summary(system_prompt, user_prompt)
                        # print(f"Finish index {index}")
                        break
                    except Exception as e:
                        if 'Rate limit reached' in str(e):
                            wait_time = 300
                            current_time = datetime.now().strftime('%H:%M:%S')
                            print(f"Rate limit hit at {current_time}. Waiting for 5 minutes before retrying...")
                            time.sleep(wait_time)
                        elif 'is currently loading' in str(e):
                            wait_time = 200
                            print(f"Model is loading, wait for {wait_time}")
                            time.sleep(wait_time)
                        elif '429' in str(e): # for gemini models
                            wait_time = 60
                            print(f"Quota has reached, wait for {wait_time}")
                            time.sleep(wait_time)
                        else:
                            print(f"Error at index {index}: {e}")
                            _summary = ""
                            exceptions.append(index)
                            break

                summary.append(_summary)
                source.append(_source)
                dataset.append(_dataset)

                # Sleep to prevent hitting rate limits too frequently
                time.sleep(1)

            self.summaries_df = pd.DataFrame(list(zip(source, summary, dataset)),
                                            columns=["source", "summary", "dataset"])

            if save_path is not None:
                print(f'Save summaries to {save_path}')
                fpath = Path(save_path)
                fpath.parent.mkdir(parents=True, exist_ok=True)
                self.summaries_df.to_csv(fpath) 

        self.exceptions = exceptions
        self._compute_avg_length()
        self._compute_answer_rate()

        return self.summaries_df
    
    def generate_summary(self, system_prompt: str, user_prompt: str):
        # Using Together AI API
        using_together_api = False
        together_ai_api_models = ['mixtral', 'dbrx', 'wizardlm', 'llama-3-', 'qwen2-72b-instruct', 'zero-one-ai', 'llama-3.2-'] #, 'mistralai'
        using_replicate_api = False
        replicate_api_models = ['snowflake', 'llama-3.1-405b']
        using_pipeline = False
        pipeline_models = ['llama-3.1', 'phi-3-mini','falcon-7b', 'phi-3.5', 'mistral-nemo', 'llama-3.3', 'phi-4']

        for replicate_api_model in replicate_api_models:
            if replicate_api_model in self.model_id.lower():
                using_replicate_api = True
                break

        if not using_replicate_api:
            for together_ai_api_model in together_ai_api_models:
                if together_ai_api_model in self.model_id.lower():
                    using_together_api = True
                    break

        if not using_replicate_api and not using_together_api:
            for pipeline_model in pipeline_models:
                if pipeline_model in self.model_id.lower():
                    using_pipeline = True
                    break

        # if 'mixtral' in self.model_id.lower() or 'dbrx' in self.model_id.lower() or 'wizardlm' in self.model_id.lower(): # For mixtral and dbrx models, use Together AI API
        if using_together_api:
            print('using together api')
            client = Together(api_key=os.environ.get('TOGETHER_API_KEY'))
            if 'llama-3.2-90b-vision' in self.model_id.lower() or 'llama-3.2-11b-vision' in self.model_id.lower():
                messages = [
                        {"role": "system","content": system_prompt},
                        {"role": "user","content": [{"type": "text","text": user_prompt}]}
                ]
            else:
                messages = [{"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}]
            response = client.chat.completions.create(
                model=self.model_id,
                messages = messages,
                max_tokens=250,
                temperature=0,
            )
            # print(response)
            result = response.choices[0].message.content
            print(result)
            return result

        # Using OpenAI API
        elif 'openai' in self.model_id.lower():
            client = OpenAI()
            response = client.chat.completions.create(
                model=self.model_id.replace('openai/',''),
                messages=[{"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}] if 'gpt' in self.model_id
                        else [{"role": "user", "content": system_prompt + '\n' + user_prompt}],
                temperature=0.0 if 'gpt' in self.model_id.lower() else 1.0, # fixed at 1 for o1 models 
                # max_completion_tokens=250 if 'gpt' in self.model_id.lower() else None, #  not compatible with o1 series models
            )   
            # print(response)
            result = response.choices[0].message.content
            print(result)
            return result

        # Using Grok API
        elif 'grok' in self.model_id.lower(): # xai
            XAI_API_KEY = os.getenv("XAI_API_KEY")
            client = OpenAI(
                api_key=XAI_API_KEY,
                base_url="https://api.x.ai/v1",
            )

            completion = client.chat.completions.create(
                model=self.model_id.split('/')[-1],
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                temperature=0.0
            )
            result = completion.choices[0].message.content
            print(result)
            return result

        # Using Vertex AI API for Gemini models
        elif 'gemini' in self.model_id.lower():
            vertexai.init(project=os.getenv("GOOGLE_PROJECT_ID"), location="us-central1")
            model = GenerativeModel(
                self.model_id.lower().split('google/')[-1],
                system_instruction = [system_prompt]
            )
            generation_config = {
                "temperature": 0,
                "max_output_tokens": 500
            }
            safety_settings = [
                SafetySetting(
                    category=SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
                    threshold=SafetySetting.HarmBlockThreshold.BLOCK_NONE
                ),
                SafetySetting(
                    category=SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
                    threshold=SafetySetting.HarmBlockThreshold.BLOCK_NONE
                ),
                SafetySetting(
                    category=SafetySetting.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
                    threshold=SafetySetting.HarmBlockThreshold.BLOCK_NONE
                ),
                SafetySetting(
                    category=SafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT,
                    threshold=SafetySetting.HarmBlockThreshold.BLOCK_NONE
                )
            ]
            response = model.generate_content(
                user_prompt,
                safety_settings=safety_settings,
                generation_config=generation_config
            )
            result = response.text
            print(result)
            return result
        
        # Using Replicate API
        elif using_replicate_api:
            print("using replicate")
            if 'snowflake' in self.model_id.lower():
                input = {
                    "prompt": user_prompt,
                    "temperature": 0,
                    "max_new_tokens": 250,
                    "stop_sequences": "<|im_end|>",
                    "prompt_template": f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + "<|im_start|>user\n{prompt}<|im_end|>\n\n<|im_start|>assistant\n",
                }
            else:
                input = {
                    "prompt": user_prompt,
                    "system_prompt": system_prompt,
                    "temperature": 0,
                    "max_new_tokens": 250
                }
            response = replicate.run(
                self.model_id,
                input=input
            )
            # print(response)
            if isinstance(response, list):
                response = ''.join(response)
                # print(response)
                # print()
            print(response)
            return response

        # Using Anthropic API for Claude models
        elif 'claude' in self.model_id.lower(): # using anthropic api
            print('using Anthropic API')
            client = anthropic.Anthropic()
            message = client.messages.create(
                model=self.model_id.split('/')[-1],
                max_tokens=1024,
                temperature=0,
                system=system_prompt,
                messages=[
                    {
                        "role": "user",
                        # "content": [
                        #     {
                        #         "type": "text",
                        #         "text": user_prompt
                        #     }
                        # ]
                        "content": user_prompt
                    }
                ]
            )
            result = message.content[0].text
            print(result)
            return result

        # Using Cohere API
        elif 'command-r' in self.model_id.lower() or 'aya-expanse' in self.model_id.lower():
            co = cohere.ClientV2(os.getenv('COHERE_API_TOKEN'))
            response = co.chat(
                model=self.model_id.split('/')[-1],
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0,
            )
            result = response.message.content[0].text
            print(result)
            return result

        # Using MistralAI API
        elif 'mistral-large' in self.model_id.lower():
            api_key = os.environ["MISTRAL_API_KEY"]
            client = Mistral(api_key=api_key)

            messages = [
                {
                    "role":"system", 
                    "content":system_prompt
                },
                {
                    "role":"user", 
                    "content":user_prompt
                }
            ]

            # No streaming
            chat_response = client.chat.complete(
                model=self.model_id,
                messages=messages,
            )
            result = chat_response.choices[0].message.content
            print(result)
            return result

        # Using Deepseek API
        elif 'deepseek' in self.model_id.lower():
            client = OpenAI(api_key=os.getenv("DeepSeek_API_KEY"), base_url="https://api.deepseek.com")
            response = client.chat.completions.create(
                model=self.model_id.split('/')[-1],
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                max_tokens=250,
                temperature=0,
                stream=False
            )
            result = response.choices[0].message.content
            print(result)
            return result
        
        # Using HF pipeline or local checkpoints
        elif self.local_model is None and self.local_pipeline is None:
            if using_pipeline:
                self.local_pipeline = pipeline(
                    "text-generation",
                    model=self.model_id,
                    tokenizer=AutoTokenizer.from_pretrained(self.model_id),
                    torch_dtype=torch.bfloat16 if 'llama-3.2' in self.model_id.lower() or 'llama-3.3' in self.model_id.lower() else "auto",
                    device_map="auto",
                    trust_remote_code=True
                )
            else:
                if 'ragamuffin' in self.model_id.lower():
                    self.tokenizer = AutoTokenizer.from_pretrained(os.path.join('/home/miaoran', self.model_id))

                else:
                    self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf" if 'openelm' in self.model_id.lower() else self.model_id, trust_remote_code=True)
                print("Tokenizer loaded")
                if 'jamba' in self.model_id.lower():
                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id,
                                             torch_dtype=torch.bfloat16,
                                             attn_implementation="flash_attention_2",
                                             device_map="auto",
                                             use_mamba_kernels=False)
                    
                elif 'qwen2-vl' in self.model_id.lower():
                    self.local_model = Qwen2VLForConditionalGeneration.from_pretrained(
                        self.model_id, torch_dtype="auto", device_map="auto"
                    )
                    self.processor = AutoProcessor.from_pretrained(self.model_id)

                # elif 'ragamuffin' in self.model_id.lower():
                #     print('Using ragamuffin')
                #     self.local_model = AutoModelForCausalLM.from_pretrained(os.path.join('/home/miaoran', self.model_id),
                #                                  torch_dtype=torch.bfloat16, # forcing bfloat16 for now
                #                                  attn_implementation="flash_attention_2")
                elif 'olmo' in self.model_id.lower():
                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id)#torch_dtype="auto"
                
                elif 'qwq-' in self.model_id.lower():
                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id, torch_dtype="auto", device_map="auto")

                else:
                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id, trust_remote_code=True, device_map="auto")#torch_dtype="auto"
                # print(self.local_model.device)
                print("Local model loaded")
                    
        # Using local model/pipeline
        if self.local_pipeline:
            print('Using Transformers pipeline')
            messages=[
                {"role": "system", "content": system_prompt}, 
                {"role": "user", "content": user_prompt}
            ]
            outputs = self.local_pipeline(
                messages,
                max_new_tokens=256,
                # return_full_text=False,
                do_sample=False
            )
            result = outputs[0]["generated_text"][-1]['content']
            print(result)
            return result

        elif self.local_model: # cannot call API. using local model / pipeline
            print('Using local model')

            # Set appropriate prompt based on model document
            if 'gemma' in self.model_id.lower() or 'mistral-7b' in self.model_id.lower():
                messages=[
                    # gemma-1.1, mistral-7b does not accept system role
                    {"role": "user", "content": system_prompt + '\n' + user_prompt}
                ]
                prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
            
            elif 'phi-2' in self.model_id.lower():
                prompt = system_prompt + '\n' + user_prompt

            elif 'intel' in self.model_id.lower():
                prompt = f"### System:\n{system_prompt}\n### User:\n{user_prompt}\n### Assistant:\n"

            elif 'qwen2-vl' in self.model_id.lower():
                messages = [
                    {   
                        "role": "system",
                        "content": [
                            {"type": "text", "text": system_prompt}
                        ]
                    }, 
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": user_prompt},
                        ],
                    }
                ]
            else:
                messages=[
                    {"role": "system", "content": system_prompt}, 
                    {"role": "user", "content": user_prompt}
                ]
                prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
           
            # Tokenize inputs
            if 'olmo' in self.model_id.lower():
                input_ids = self.tokenizer([prompt], return_tensors='pt', return_token_type_ids=False)#.to(self.device)
            elif 'qwq' in self.model_id.lower():
                input_ids = self.tokenizer([prompt], return_tensors="pt").to(self.device)
            else:
                input_ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            
            # Generate outputs
            if 'granite' in self.model_id.lower():
                self.local_model.eval()
                outputs = self.local_model.generate(**input_ids, max_new_tokens=250)
            elif 'olmo' in self.model_id.lower():
                outputs = self.local_model.generate(**input_ids, max_new_tokens=250, do_sample=True, temperature=0.01)#top_k=50, top_p=0.95)
            elif 'qwq' in self.model_id.lower():
                outputs = self.local_model.generate(**input_ids, max_new_tokens=512, do_sample=True, temperature=0.01)
            else:
                with torch.no_grad():
                    outputs = self.local_model.generate(**input_ids, do_sample=True, max_new_tokens=250, temperature=0.01)#, pad_token_id=self.tokenizer.eos_token_id
            if 'glm' in self.model_id.lower() or 'ragamuffin' in self.model_id.lower() or 'granite' in self.model_id.lower():
                outputs = outputs[:, input_ids['input_ids'].shape[1]:]
            elif 'qwen2-vl' in self.model_id.lower() or 'qwen2.5' in self.model_id.lower() or 'qwq-' in self.model_id.lower():
                outputs = [
                    out_ids[len(in_ids) :] for in_ids, out_ids in zip(input_ids.input_ids, outputs)
                ]
            
            # Decode outputs
            if 'qwen2-vl' in self.model_id.lower():
                result = self.processor.batch_decode(
                    outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False
                )[0]
            elif 'olmo' in self.model_id.lower() or 'qwq' in self.model_id.lower():
                result = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
            else:
                result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

            if 'gemma-2' in self.model_id.lower():
                result = result.split(user_prompt + '\nmodel')[-1].strip()
            elif 'intel' in self.model_id.lower():
                result = result.split("### Assistant:\n")[-1]
            elif 'jamba' in self.model_id.lower():
                result = result.split(messages[-1]['content'])[1].strip()
            elif 'qwen2-vl' in self.model_id.lower() or 'qwen2.5' in self.model_id.lower():
                pass
            elif 'olmo' in self.model_id.lower():
                result = result.split("<|assistant|>\n")[-1]
            else:
                result = result.replace(prompt.strip(), '')
            
            print(result)
            return result

    def _compute_avg_length(self):
        """
        Compute the average length of non-empty summaries using SpaCy.
        """
        total_word_count = 0
        total_count = 0

        for summary in self.summaries_df['summary']:
            if util.is_summary_valid(summary):
                doc = nlp(summary)
                words = [token.text for token in doc if token.is_alpha]
                total_word_count += len(words)
                total_count += 1

        self.avg_length = 0 if total_count == 0 else total_word_count / total_count

    def _compute_answer_rate(self):
        """
        Compute the rate of non-empty summaries.
        """
        valid_count = sum(1 for summary in self.summaries_df['summary']
                            if util.is_summary_valid(summary))

        total_count = len(self.summaries_df)

        self.answer_rate = 0 if total_count == 0 else valid_count / total_count


class EvaluationModel:
    """A class to evaluate generated summaries.

    Attributes:
        model (CrossEncoder): The evaluation model.
        scores (list): List of evaluation scores.
        accuracy (float): Accuracy of the summaries.
        hallucination_rate (float): Rate of hallucination in summaries.
    """

    def __init__(self, model_path, device):
        """
        Initializes the EvaluationModel with a CrossEncoder model.

        Args:
            model_path (str): Path to the CrossEncoder model.
        """
        config = AutoConfig.from_pretrained('google/flan-t5-large')
        self.model = AutoModelForTokenClassification.from_pretrained(model_path, config=config)
        self.device = device
        self.model.to(self.device)
        self.scores = []
        self.factual_consistency_rate = None
        self.hallucination_rate = None
    
    def predict(self, text_pairs):
        """Load LoRA adapters of HHEM and make predictions
        All HHEM 2.1 settings, e.g., prompt template, are hardcoded in this function.
        Args:
            text_pairs: list of tuples, each tuple contains two strings (premise, hypothesis)
            checkpoint: model ID on Hugging Face
        """

        prompt = "<pad> Determine if the hypothesis is true given the premise?\n\nPremise: {text1}\n\nHypothesis: {text2}"
        
        tokenizer = AutoTokenizer.from_pretrained('t5-base')
        inputs = tokenizer(
            [prompt.format(text1=pair[0], text2=pair[1]) for pair in text_pairs], 
            return_tensors='pt', padding='longest').to(self.device)
        
        self.model.eval()
        with torch.no_grad():
            output = self.model(**inputs)
        logits = output.logits
        logits = logits[:,0,:] # get the logits on the first token
        logits = torch.softmax(logits, dim=-1)
        scores = [round(x, 5) for x in logits[:, 1].tolist()] # list of float
        return scores

    def evaluate_hallucination(self, summaries_df):
        """
        Evaluate the hallucination rate in summaries. Updates the 'scores' attribute 
        of the instance with the computed scores.

        Args:
            summaries_df (DataFrame): DataFrame containing source docs and summaries.

        Returns:
            list: List of hallucination scores. Also updates the 'scores' attribute of the instance.
        """
        hem_scores = []
        sources = []
        summaries = []
        source_summary_pairs = util.create_pairs(summaries_df)

        for doc, summary in source_summary_pairs:
            if util.is_summary_valid(summary):
                try:
                    summary = util.normalize_summary(summary)
                    score = self.predict([(doc, summary)])[0]
                    hem_scores.append(score)
                    sources.append(doc)
                    summaries.append(summary)
                    if score < 0.5:
                        print(score)
                        print(doc)
                        print('-'*20)
                        print(summary)
                        print('='*50)
                except Exception as e:
                    logging.error(f"Error while running HEM: {e}")
                    raise

        self.scores = hem_scores
        eval_results = {'source': sources, 'summary': summaries, 'HEM scores': hem_scores}
        return hem_scores, eval_results


    def compute_factual_consistency_rate(self, threshold=0.5):
        """
        Compute the factual consistency rate of the evaluated summaries based on
        the previously calculated scores. This method relies on the 'scores'
        attribute being populated, typically via the 'evaluate_hallucination' method.

        Returns:
            float: Factual Consistency Rate. Also updates the 'factual_consistency_rate'
            and 'hallucination_rate' attributes of the instance.

        Raises:
            ValueError: If scores have not been calculated prior to calling this method.
        """
        if not self.scores:
            error_msg = "Scores not calculated. Call evaluate_hallucination() first."
            logging.error(error_msg)
            raise ValueError(error_msg)

        # Use threshold of 0.5 to compute factual_consistency_rate
        num_above_threshold = sum(score >= threshold for score in self.scores)
        num_total = len(self.scores)

        if not num_total:
            raise ValueError("No scores available to compute factual consistency rate.")

        self.factual_consistency_rate = (num_above_threshold / num_total) * 100
        self.hallucination_rate = 100 - self.factual_consistency_rate

        return self.factual_consistency_rate