minor fix

src/backend/model_operations.py

@@ -13,7 +13,7 @@ from sentence_transformers import CrossEncoder
 import litellm
 # from litellm import completion
 from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig, pipeline
 # from accelerate import PartialState
 # from accelerate.inference import prepare_pippy
 import torch
@@ -272,9 +272,8 @@ class SummaryGenerator:
         # Using HF API or download checkpoints
         elif self.local_model is None:
             try: # try use HuggingFace API
-                
                 response = litellm.completion(
-                    model='command-r-plus' if 'command' in self.model else self.model,
+                    model='command-r-plus' if 'command' in self.model_id else self.model_id,
                     messages=[{"role": "system", "content": system_prompt},
                                 {"role": "user", "content": user_prompt}],
                     temperature=0.0,
@@ -286,7 +285,7 @@ class SummaryGenerator:
             except: # fail to call api. run it locally.
                 self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
                 print("Tokenizer loaded")
-                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", cache_dir='/home/paperspace/cache')
                 print("Local model loaded")
         
         # Using local model
@@ -294,15 +293,33 @@ class SummaryGenerator:
             messages=[
                 {"role": "system", "content": system_prompt}, # gemma-1.1 does not accept system role
                 {"role": "user", "content": user_prompt}
-            ],
-            prompt = self.tokenizer.apply_chat_template(messages,add_generation_prompt=True, tokenize=False)
-            print(prompt)
-            input_ids = self.tokenizer(prompt, return_tensors="pt").to('cuda')
-            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)
-            result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            result = result.replace(prompt[0], '')
-            print(result)
+            ]
+            try: # some models support pipeline
+                pipe = pipeline(
+                    "text-generation",
+                    model=self.local_model,
+                    tokenizer=self.tokenizer,
+                )
+
+                generation_args = {
+                    "max_new_tokens": 250,
+                    "return_full_text": False,
+                    "temperature": 0.0,
+                    "do_sample": False,
+                }
+
+                output = pipe(messages, **generation_args)
+                result = output[0]['generated_text']
+                print(result)
+            except:
+                prompt = self.tokenizer.apply_chat_template(messages,add_generation_prompt=True, tokenize=False)
+                print(prompt)
+                input_ids = self.tokenizer(prompt, return_tensors="pt").to('cuda')
+                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)
+                result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+                result = result.replace(prompt[0], '')
+                print(result)
             return result
 
     def _compute_avg_length(self):