new attempt on zero

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import spaces  # Import the spaces module for ZeroGPU compatibility
 
 # Load the model and tokenizer from Hugging Face
 model_id = "seedboxai/KafkaLM-8x7B-German-V0.1-DPO"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True, load_in_8bit=False, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True, load_in_8bit=False, trust_remote_code=True).to('cuda')  # Move the model to GPU
 
 # Define the text generation function
+@spaces.GPU  # Decorate this function to use GPU
 def generate_text(user_input, system_prompt):
     # Combine the system prompt and the user input to form the full prompt
     full_prompt = f"{system_prompt.strip()}\n\n{user_input.strip()}"
@@ -14,7 +16,7 @@ def generate_text(user_input, system_prompt):
     # Initialize the pipeline for text generation
     text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer,
                               return_full_text=True, temperature=0.5,
-                              max_new_tokens=512, top_p=0.95, top_k=50, do_sample=True)
+                              max_new_tokens=512, top_p=0.95, top_k=50, do_sample=True, device=0)  # Specify the device for the pipeline
     
     # Generate text based on the full prompt
     results = text_generator(full_prompt)
@@ -24,7 +26,7 @@ def generate_text(user_input, system_prompt):
 
 # Setup the Gradio interface
 iface = gr.Interface(fn=generate_text,
-                     inputs=[gr.Textbox(lines=2, label="User Prompt"), gr.Textbox(lines=5, label="System Prompt",text="Du bist ein freundlicher und hilfsbereiter KI-Assistent. Du beantwortest Fragen faktenorientiert und präzise, ohne dabei relevante Fakten auszulassen.")],
+                     inputs=[gr.Textbox(lines=2, label="User Prompt"), gr.Textbox(lines=5, label="System Prompt", text="Du bist ein freundlicher und hilfsbereiter KI-Assistent. Du beantwortest Fragen faktenorientiert und präzise, ohne dabei relevante Fakten auszulassen.")],
                      outputs=gr.Textbox(label="Generated Text"),
                      title="Text Generation with KafkaLM",
                      description="Enter a user prompt and a system prompt to generate text using the KafkaLM model.")