Update app.py

app.py

@@ -17,18 +17,24 @@ print('token = ',token)
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_id = "mistralai/Mistral-7B-v0.3"
+# model_id = "mistralai/Mistral-7B-v0.3"
 
-model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+# model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
+from airllm import AirLLMLlama2
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, token= token)
+MAX_LENGTH = 128
+# could use hugging face model repo id:
+model = AirLLMLlama2("garage-bAInd/Platypus2-70B-instruct")
 
-model = AutoModelForCausalLM.from_pretrained(model_id, token= token, torch_dtype=torch.bfloat16,    
-                                             # attn_implementation="flash_attention_2",
-                                             # low_cpu_mem_usage=True,
-                                             device_map="auto"
-                                            )
+
+# tokenizer = AutoTokenizer.from_pretrained(model_id, token= token)
+
+# model = AutoModelForCausalLM.from_pretrained(model_id, token= token, torch_dtype=torch.bfloat16,    
+#                                              # attn_implementation="flash_attention_2",
+#                                              # low_cpu_mem_usage=True,
+#                                              device_map="auto"
+#                                             )
 
 
 
@@ -41,6 +47,27 @@ def respond(
     temperature,
     top_p,
 ):
+    input_text = [
+        'What is the capital of United States?',
+    ]
+
+    input_tokens = model.tokenizer(input_text,
+    return_tensors="pt", 
+    return_attention_mask=False, 
+    truncation=True, 
+    max_length=MAX_LENGTH, 
+    padding=True)
+           
+    generation_output = model.generate(
+    input_tokens['input_ids'].cuda(), 
+    max_new_tokens=20,
+    use_cache=True,
+    return_dict_in_generate=True)
+
+    output = model.tokenizer.decode(generation_output.sequences[0])
+
+    print(output)
+    yield output
     
   
     messages = [
@@ -49,13 +76,13 @@ def respond(
     {"role": "user", "content": "Do you have mayonnaise recipes?"}
 ]
 
-    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")
+    # inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")
 
-    outputs = model.generate(inputs, max_new_tokens=2000)
-    gen_text=tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # outputs = model.generate(inputs, max_new_tokens=2000)
+    # gen_text=tokenizer.decode(outputs[0], skip_special_tokens=True)
    
-    print(gen_text)
-    yield gen_text
+    # print(gen_text)
+    # yield gen_text
     # for val in history:
     #     if val[0]:
     #         messages.append({"role": "user", "content": val[0]})