minor fix to work for recent models

src/backend/model_operations.py

@@ -11,15 +11,19 @@ import pandas as pd
 import spacy
 import litellm
 from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, AutoModelForTokenClassification, AutoConfig
+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 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
@@ -156,7 +160,7 @@ class SummaryGenerator:
     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-', 'qwen', 'zero-one-ai'] #, 'mistralai'
+        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
@@ -181,99 +185,80 @@ class SummaryGenerator:
 
         # 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')
-            # suffix = "completions" if ('mixtral' in self.model_id.lower() or 'base' in self.model_id.lower()) else "chat/completions"
-            suffix = "chat/completions"
-            url = f"https://api.together.xyz/v1/{suffix}"
-
-            payload = {
-                "model": self.model_id,
-                'max_new_tokens': 250,
-                "temperature": 0.0,
-        
-            }
-            payload['messages'] = [{"role": "system", "content": system_prompt},
-                                        {"role": "user", "content": user_prompt}]
-            headers = {
-                "accept": "application/json",
-                "content-type": "application/json",
-                "Authorization": f"Bearer {os.environ['TOGETHER_API_KEY']}"
-            }
-
-            response = requests.post(url, json=payload, headers=headers)
-            print(response)
-            try:
-                result = json.loads(response.text)
-                # print(result)
-                result = result["choices"][0]
-                if 'message' in result:
-                    result = result["message"]["content"].strip()
-                else:
-                    result = result["text"]
-                    result_candidates = [result_cancdidate for result_cancdidate in result.split('\n\n') if len(result_cancdidate) > 0]
-                    result = result_candidates[0]
-                # print(result)
-            except:
-                # print(response)
-                result = ''
+            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 'gpt' in self.model_id.lower():
+        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}],
-                temperature=0.0,
-                max_tokens=250,
+                        {"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 Google AI API for Gemini models
+
         elif 'gemini' in self.model_id.lower():
-            genai.configure(api_key=os.getenv('GOOGLE_AI_API_KEY'))
+            vertexai.init(project=os.getenv("GOOGLE_PROJECT_ID"), location="us-central1")
+            gemini_model_id_map = {'gemini-1.5-pro-exp-0827':'gemini-pro-experimental', 'gemini-1.5-flash-exp-0827': 'gemini-flash-experimental'}
+            model = GenerativeModel(
+                self.model_id.lower().split('google/')[-1],
+                system_instruction = [system_prompt]
+            )
             generation_config = {
                 "temperature": 0,
-                "top_p": 0.95, # cannot change
-                "top_k": 0,
-                "max_output_tokens": 250,
-                # "response_mime_type": "application/json",
+                "max_output_tokens": 250
             }
             safety_settings = [
-                {
-                    "category": "HARM_CATEGORY_HARASSMENT",
-                    "threshold": "BLOCK_NONE"
-                },
-                {
-                    "category": "HARM_CATEGORY_HATE_SPEECH",
-                    "threshold": "BLOCK_NONE"
-                },
-                {
-                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-                    "threshold": "BLOCK_NONE"
-                },
-                {
-                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-                    "threshold": "BLOCK_NONE"
-                },
+                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
+                )
             ]
-            model = genai.GenerativeModel(model_name=self.model_id.lower().split('google/')[-1],
-                              generation_config=generation_config,
-                              system_instruction=system_prompt,
-                              safety_settings=safety_settings)
-            # print(model)
-            convo = model.start_chat(history=[])
-            convo.send_message(user_prompt)
-            # print(convo.last)
-            result = convo.last.text
+            response = model.generate_content(
+                user_prompt,
+                safety_settings=safety_settings,
+                generation_config=generation_config
+            )
+            result = response.text
             print(result)
             return result
-
         elif using_replicate_api:
             print("using replicate")
             if 'snowflake' in self.model_id.lower():
@@ -338,7 +323,6 @@ class SummaryGenerator:
             print(result)
             return result
 
-
         elif 'mistral-large' in self.model_id.lower():
             api_key = os.environ["MISTRAL_API_KEY"]
             client = Mistral(api_key=api_key)
@@ -363,35 +347,30 @@ class SummaryGenerator:
             print(result)
             return result
 
+        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 API or download checkpoints
         elif self.local_model is None and self.local_pipeline is None:
-            # try: # try use HuggingFace API
-            #     print('** using huggingface api')
-            #     response = litellm.completion(
-            #         model=self.model,
-            #         messages=[{"role": "system", "content": system_prompt},
-            #                     {"role": "user", "content": user_prompt}],
-            #         temperature=0.0,
-            #         max_tokens=250,
-            #         api_base=self.api_base,
-            #     )
-            #     result = response['choices'][0]['message']['content']
-            #     result = result.split('<|im_end|>')[0]
-            #     print(result)
-            #     return result
-            # 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)
-            #     else:
             if using_pipeline:
                 self.local_pipeline = pipeline(
                     "text-generation",
                     model=self.model_id,
                     tokenizer=AutoTokenizer.from_pretrained(self.model_id),
-                    model_kwargs={"torch_dtype": torch.bfloat16},
+                    torch_dtype=torch.bfloat16 if 'llama-3.2' in self.model_id.lower() else "auto",
                     device_map="auto",
                     trust_remote_code=True
                 )
@@ -404,12 +383,18 @@ class SummaryGenerator:
                                              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)
+
                 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')
@@ -432,7 +417,7 @@ class SummaryGenerator:
             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 + ' ' + user_prompt}
+                    {"role": "user", "content": system_prompt + '\n' + user_prompt}
                 ]
                 prompt = self.tokenizer.apply_chat_template(messages,add_generation_prompt=True, tokenize=False)
             
@@ -442,6 +427,21 @@ class SummaryGenerator:
             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}, 
@@ -455,14 +455,27 @@ class SummaryGenerator:
                 outputs = self.local_model.generate(**input_ids, max_new_tokens=250, do_sample=True, temperature=0.01, pad_token_id=self.tokenizer.eos_token_id)
             if 'glm' in self.model_id.lower():
                 outputs = outputs[:, input_ids['input_ids'].shape[1]:]
-                
-            result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            elif 'qwen2-vl' in self.model_id.lower() or 'qwen2.5' in self.model_id.lower():
+                outputs = [
+                    out_ids[len(in_ids) :] for in_ids, out_ids in zip(input_ids.input_ids, 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]
+            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
             else:
                 # print(prompt)
                 # print('-'*50)
@@ -572,6 +585,12 @@ class EvaluationModel:
                     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