switched to pipeline

app.py

@@ -7,19 +7,35 @@ import transformers
 # Load the model and tokenizer
 model_name = "Artigenz/Artigenz-Coder-DS-6.7B"
 tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
-model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
+max_new_tokens:int=1024
+do_sample:bool=True
+num_beams:int=1
+temperature:float=0.5
+top_p:float=0.95
+top_k:float=40
+repetition_penalty:float=1.1
+pipe = transformers.pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        num_beams=num_beams,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+    )
 
-# Function to generate responses from the model
 def generate_response(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt")
-    input_ids = inputs["input_ids"].to(model.device)
-    attention_mask = inputs["attention_mask"].to(model.device)
-
-    # Generate the output
-    outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=512)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return response
+    messages = [
+                {
+                    "role": "system", "content": "You are a helpful coding chatbot. You will answer the user's questions to the best of your ability.",
+                    "role": "user", "content": input_text,
+                },
+        ]
+    return pipe(messages)[0]['generated_text'][-1]['content'].replace("\\n", "\n")
 
 # Define the Gradio interface
 iface = gr.Interface(