Update app.py

app.py

@@ -36,13 +36,20 @@ def generate_text(input: PromptInput):
     prompt = input.prompt  # Access prompt from the request body
 
     # Format the prompt with specific style for your fine-tuned model
-    inputs = tokenizer(f"### Prompt: {prompt}\n### Completion:", return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    input_text = f"### Prompt: {prompt}\n### Completion:"
+    inputs = tokenizer(input_text, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
 
     # Generate the output
     with torch.no_grad():
-        output = model.generate(**inputs, max_length=128)
+        output = model.generate(**inputs, max_length=128, do_sample=False, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id)
     
     # Decode the output and remove special tokens
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    return {"generated_query": generated_text}
+    # Extract the query part from the generated output
+    if "### Completion:" in generated_text:
+        query_output = generated_text.split("### Completion:")[1].strip()
+    else:
+        query_output = generated_text.replace(input_text, "").strip()  # Fallback if the structure is not as expected
+
+    return {"generated_query": query_output}