update summary generation for new models

src/backend/model_operations.py

@@ -215,16 +215,34 @@ class SummaryGenerator:
                         {"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
+                # 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
 
+        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
+
         elif 'gemini' in self.model_id.lower():
             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]
@@ -289,21 +307,23 @@ class SummaryGenerator:
             return response
 
         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=250,
+                max_tokens=1024,
                 temperature=0,
                 system=system_prompt,
                 messages=[
                     {
                         "role": "user",
-                        "content": [
-                            {
-                                "type": "text",
-                                "text": user_prompt
-                            }
-                        ]
+                        # "content": [
+                        #     {
+                        #         "type": "text",
+                        #         "text": user_prompt
+                        #     }
+                        # ]
+                        "content": user_prompt
                     }
                 ]
             )
@@ -311,15 +331,17 @@ class SummaryGenerator:
             print(result)
             return result
 
-        elif 'command-r' in self.model_id.lower():
-            co = cohere.Client(os.getenv('COHERE_API_TOKEN'))
+        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(
-                chat_history=[
-                    {"role": "SYSTEM", "message": system_prompt},
+                model=self.model_id.split('/')[-1],
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt}
                 ],
-                message=user_prompt,
+                temperature=0,
             )
-            result = response.text
+            result = response.message.content[0].text
             print(result)
             return result
 
@@ -375,7 +397,10 @@ class SummaryGenerator:
                     trust_remote_code=True
                 )
             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)
+                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,
@@ -390,8 +415,14 @@ class SummaryGenerator:
                     )
                     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")
+    
                 else:
-                    self.local_model = AutoModelForCausalLM.from_pretrained(self.model_id, trust_remote_code=True, device_map="auto", torch_dtype="auto")
+                    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")
                     
@@ -419,7 +450,7 @@ class SummaryGenerator:
                     # 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)
+                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
@@ -451,20 +482,25 @@ class SummaryGenerator:
             # print(prompt)
             # print('-'*50)
             input_ids = self.tokenizer(prompt, return_tensors="pt").to(self.device)
-            with torch.no_grad():
-                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():
+            if 'granite' in self.model_id.lower():
+                self.local_model.eval()
+                outputs = self.local_model.generate(**input_ids, max_new_tokens=250)
+            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():
                 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]
+            # elif 'granite' in self.model_id.lower():
+            #     result = self.tokenizer.batch_decode(outputs)[0]
             else:
                 result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)