Update app.py

app.py

@@ -8,27 +8,15 @@ model = AutoModelForSeq2SeqLM.from_pretrained(model_repo_path)
 tokenizer = AutoTokenizer.from_pretrained(model_repo_path)
 
 # Function to generate Python code from LaTeX expression
-def generate_code(latex_expr, tokenizer, model, max_length=512):
-    # Tokenize the input LaTeX expression
-    inputs = tokenizer.encode(
-        latex_expr,
-        return_tensors='pt',
-        max_length=max_length,
-        truncation=True
-    )
-
+def generate_code_from_latex(latex_expression, max_length=256):
+    inputs = tokenizer(f"Latex Expression: {latex_expression} Solution:", return_tensors="pt").to(model.device)
+    
     # Generate the output
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs,
-            max_length=max_length,
-            num_beams=5,  # Number of beams for beam search
-            early_stopping=True
-        )
-
-    # Decode the output tokens to string
+    outputs = model.generate(**inputs, max_length=max_length)
+    
+    # Decode the output into Python code
     generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
+    
     return generated_code
 
 # Streamlit app layout
@@ -47,7 +35,7 @@ if st.button("Generate Code"):
         st.session_state.latex_expr = latex_input
         with st.spinner("Generating Python Code..."):
             try:
-                generated_code = generate_code(latex_expr=st.session_state.latex_expr, tokenizer=tokenizer, model=model)
+                generated_code = generate_code(latex_expr=st.session_state.latex_expr)
                 # Display the generated code
                 st.subheader("Generated Python Code")
                 st.code(generated_code, language='python')