Add cuda and sampling pram

app.py

@@ -4,40 +4,51 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from huggingface_hub import login
 
-login(token=os.getenv('HF_TOKEN'))
-
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-7B")
 model = AutoModelForCausalLM.from_pretrained(
     "Zyphra/Zamba2-7B",
-    device_map="auto",  # Automatically handles device placement
+    device_map="cuda",  # Automatically handles device placement
     torch_dtype=torch.bfloat16
 )
 
-def generate_response(input_text):
-    input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
+# Define the function to generate responses
+def generate_response(input_text, max_new_tokens, temperature, top_k, top_p, repetition_penalty, num_beams, length_penalty):
+    # Tokenize and move input to model's device
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)
+    
+    # Generate response using specified parameters
     outputs = model.generate(
-        **input_ids,
-        max_new_tokens=500,
+        input_ids=input_ids,
+        max_new_tokens=max_new_tokens,
         do_sample=True,
-        temperature=0.7,
-        top_k=50,
-        top_p=0.9,
-        repetition_penalty=1.2,
-        num_beams=5,
-        length_penalty=1.0,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        num_beams=num_beams,
+        length_penalty=length_penalty,
         num_return_sequences=1
     )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
 
-# Create the Gradio interface
+# Create Gradio interface with adjustable parameters
 demo = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(lines=5, placeholder="Enter your question here..."),
-    outputs=gr.Textbox(),
+    inputs=[
+        gr.Textbox(lines=1, placeholder="Enter a text to prepend...", label="Input Text"),
+        gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens"),
+        gr.Slider(0.1, 1.5, step=0.1, value=0.7, label="Temperature"),
+        gr.Slider(1, 100, step=1, value=50, label="Top K"),
+        gr.Slider(0.1, 1.0, step=0.1, value=0.9, label="Top P"),
+        gr.Slider(1.0, 2.0, step=0.1, value=1.2, label="Repetition Penalty"),
+        gr.Slider(1, 10, step=1, value=5, label="Number of Beams"),
+        gr.Slider(0.0, 2.0, step=0.1, value=1.0, label="Length Penalty")
+    ],
+    outputs=gr.Textbox(label="Generated Response"),
     title="Zamba2-7B Model",
-    description="Ask Zamba2 7B a question."
+    description="Ask Zamba2 7B a question with customizable parameters."
 )
 
 if __name__ == "__main__":