Update app.py

app.py

@@ -4,84 +4,58 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
 
+text_generator = None
+is_hugging_face = False
+def init():
+    global text_generator
+    huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+    if not huggingface_token:
+        pass
+        print("no HUGGINGFACE_TOKEN if you need set secret ")
+        #raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
+    
+    model_id = "google/gemma-2-9b-it"
+    model_id = "Qwen/Qwen2-0.5B-Instruct"
+    
+    device = "auto" # torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = "cuda"
+    dtype = torch.bfloat16
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=huggingface_token)
+    
+    print(model_id,device,dtype)
+    histories = []
+    #model = None
 
-huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-if not huggingface_token:
-    pass
-    print("no HUGGINGFACE_TOKEN if you need set secret ")
-    #raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
-
-#model_id = "Qwen/Qwen1.5-0.5B-Chat"
-model_id = "Kendamarron/Tokara-0.5B-Chat-v0.1"
-model_id = "Qwen/Qwen2-0.5B-Instruct"
-
-device = "auto" # torch.device("cuda" if torch.cuda.is_available() else "cpu")
-dtype = torch.bfloat16
-
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=huggingface_token)
-
-print(model_id,device,dtype)
-histories = []
-#model = None
-
-
-def call_generate_text(prompt, system_message="You are a helpful assistant."):
-    if prompt =="":
-        print("empty prompt return")
-        return ""
-
-    global histories
-    #global model
-    #if model != None:# and model.is_cuda:
-    #    print("Model is alive")
-    #else:
-    #    model = AutoModelForCausalLM.from_pretrained(
-    #    model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
-    #)
+    
+    
+    if not is_hugging_face:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
+        )
+        text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device ) #pipeline has not to(device)
         
-    messages = [
-        {"role": "system", "content": system_message},
-    ]
+        if next(model.parameters()).is_cuda:
+            print("The model is on a GPU")
+        else:
+            print("The model is on a CPU")
     
-    messages += histories 
-
-    user_message = {"role": "user", "content": prompt}
-
-    messages += [user_message]
+        #print(f"text_generator.device='{text_generator.device}")
+        if str(text_generator.device).strip() == 'cuda':
+            print("The pipeline is using a GPU")
+        else:
+            print("The pipeline is using a CPU")
     
-    try:
-        text = generate_text(messages)
-        histories += [user_message,{"role": "assistant", "content": text}]
-        #model.to("cpu")
-        return text
-    except RuntimeError  as e:
-        print(f"An unexpected error occurred: {e}")
-        #model = None
-
-    return ""
-
-iface = gr.Interface(
-    fn=call_generate_text,
-    inputs=[
-        gr.Textbox(lines=3, label="Input Prompt"),
-        gr.Textbox(lines=2, label="System Message", value="あなたは親切なアシスタントで常に日本語で返答します。"),
-    ],
-    outputs=gr.Textbox(label="Generated Text"),
-    title=f"{model_id}",
-    description=f"{model_id} CPU",
-)
-print("Initialized")
+    print("initialized")
 
 @spaces.GPU(duration=120)
 def generate_text(messages):
-
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
-    )
-        
-
-    text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device) #pipeline has not to(device)
-    result = text_generator(messages, max_new_tokens=256, do_sample=True, temperature=0.7,repetition_penalty=1.1,top_p=0.95,top_k=40)
+    if is_hugging_face:#need everytime initialize for ZeroGPU
+        model = AutoModelForCausalLM.from_pretrained(
+                model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
+            )
+        text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device ) #pipeline has not to(device)
+    result = text_generator(messages, max_new_tokens=256, do_sample=True, temperature=0.7)
 
     generated_output = result[0]["generated_text"]
     if isinstance(generated_output, list):
@@ -94,6 +68,24 @@ def generate_text(messages):
     else:
         return "Unexpected output format."
 
+
+
+def call_generate_text(message, history):
+   # history.append({"role": "user", "content": message})
+    print(message)
+    print(history)
+   
+    messages = history+[{"role":"user","content":message}]
+    try:
+        text = generate_text(messages)
+        return text
+    except RuntimeError  as e:
+        print(f"An unexpected error occurred: {e}")
+       
+    return ""
+
+demo = gr.ChatInterface(call_generate_text,type="messages")
+
 if __name__ == "__main__":
-    print("Main")
-    iface.launch()
+    init()
+    demo.launch(share=True)