Update app.py

app.py

@@ -5,13 +5,47 @@ from transformers import T5ForConditionalGeneration, RobertaTokenizer
 quantized_model = T5ForConditionalGeneration.from_pretrained("AbdulHadi806/codet5-finetuned-latest-quantized")
 tokenizer = RobertaTokenizer.from_pretrained("AbdulHadi806/codet5-finetuned-latest-quantized")
 
-def inference(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = quantized_model.generate(**inputs)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def generate_code(input_text):
+    print(input_text)
+    input_ids = tokenizer(input_text, return_tensors='pt', padding="max_length", truncation=True, max_length=128).input_ids.to(model.device)
+    outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True)
+    predicted_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    cleaned_code = clean_generated_code(postprocess_output(predicted_text))
+
+    return cleaned_code
+    
+def preprocess_infer_input(text):
+    # Assuming the input is already a string, we don't need to access it as a dictionary
+    return f"latex: {text}"
+
+def clean_generated_code(generated_code):
+    # Remove unwanted parts
+    print(':::generated_code::::', generated_code)
+    cleaned_code = generated_code.replace('*convert(latex, python.code)', '').strip()
+
+    # Optionally, format the code for better readability
+    cleaned_code = cleaned_code.replace('\n', '\n').replace('  ', ' ')  # Adjust spacing if needed
+
+    return cleaned_code
+
+def generate_solution(input_text):
+    input_text = preprocess_infer_input(input_text)
+    print(input_text)
+    
+    input_ids = tokenizer(input_text, return_tensors='pt', padding="max_length", truncation=True, max_length=128).input_ids
+    input_ids = input_ids.to(model.device)
+
+    with torch.no_grad():
+        outputs = quantized_model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True)
+
+    predicted_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    cleaned_code = clean_generated_code(postprocess_output(predicted_text))
+    return cleaned_code
+
+
 
 # Create Gradio interface
-iface = gr.Interface(fn=inference, inputs="text", outputs="text")
+iface = gr.Interface(fn=generate_solution, inputs="text", outputs="text")
 
 # Launch the interface
 iface.launch()