Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from transformers import pipeline
 import logging
 import torch
+import numpy as np
 
 # Logging Setup
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -20,9 +21,38 @@ def get_model_pipeline(model_name):
         logging.error(f"Error loading model pipeline: {e}")
         return None
 
+# Beam search implementation
+def beam_search(model, prompt, beam_width=3, max_length=20):
+    sequences = [[list(prompt), 0.0]]
+    
+    for _ in range(max_length):
+        all_candidates = list()
+        
+        for seq, score in sequences:
+            if len(seq) > 0 and seq[-1] == model.tokenizer.eos_token_id:
+                all_candidates.append((seq, score))
+                continue
+            
+            inputs = model.tokenizer(seq, return_tensors='pt')
+            outputs = model.model(**inputs)
+            logits = outputs.logits[0, -1, :]
+            probabilities = torch.nn.functional.softmax(logits, dim=-1).detach().cpu().numpy()
+            
+            candidates = np.argsort(probabilities)[-beam_width:]
+            
+            for candidate in candidates:
+                new_seq = seq + [candidate]
+                new_score = score + np.log(probabilities[candidate])
+                all_candidates.append((new_seq, new_score))
+        
+        ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True)
+        sequences = ordered[:beam_width]
+        
+    return sequences[0][0]
+
 # Function to generate code
 @st.cache_data
-def generate_code(task_description, max_length, temperature, num_return_sequences, model_name):
+def generate_code(task_description, max_length, temperature, num_return_sequences, model_name, beam_width=3):
     code_pipeline = get_model_pipeline(model_name)
     if code_pipeline is None:
         return ["Error: Failed to load model pipeline."]
@@ -30,14 +60,17 @@ def generate_code(task_description, max_length, temperature, num_return_sequence
     try:
         logging.info(f"Generating code with input: {task_description}")
         prompt = f"Develop code for the following task: {task_description}"
-
-        outputs = code_pipeline(
-            prompt, 
-            max_length=max_length, 
-            num_return_sequences=num_return_sequences, 
-            temperature=temperature,
-            truncation=True  # Added truncation
-        )
+        
+        # Tokenize prompt for beam search
+        inputs = code_pipeline.tokenizer(prompt, return_tensors='pt')
+        input_ids = inputs['input_ids'][0].tolist()
+        
+        outputs = []
+        for _ in range(num_return_sequences):
+            output_tokens = beam_search(code_pipeline, input_ids, beam_width=beam_width, max_length=max_length)
+            output_text = code_pipeline.tokenizer.decode(output_tokens, skip_special_tokens=True)
+            outputs.append({'generated_text': output_text})
+        
         codes = [output['generated_text'] for output in outputs]
 
         logging.info("Code generation completed successfully.")
@@ -78,13 +111,15 @@ def main():
 
     # Options Section
     st.header("Options")
-    col1, col2, col3 = st.columns(3)
+    col1, col2, col3, col4 = st.columns(4)
     with col1:
         max_length = st.slider("Max Length", min_value=50, max_value=2048, value=250, step=50, help="Maximum length of the generated code.")
     with col2:
         temperature = st.slider("Temperature", min_value=0.1, max_value=1.0, value=0.7, step=0.1, help="Controls the creativity of the generated code.")
     with col3:
         num_return_sequences = st.slider("Number of Sequences", min_value=1, max_value=5, value=1, step=1, help="Number of code snippets to generate.")
+    with col4:
+        beam_width = st.slider("Beam Width", min_value=1, max_value=10, value=3, step=1, help="Beam width for beam search.")
 
     # Generate Code Button
     if st.button("Generate Code"):
@@ -92,7 +127,7 @@ def main():
             # Clear previous generated codes
             st.session_state.generated_codes = []
             with st.spinner("Generating code..."):
-                st.session_state.generated_codes = generate_code(task_description, max_length, temperature, num_return_sequences, model_name)
+                st.session_state.generated_codes = generate_code(task_description, max_length, temperature, num_return_sequences, model_name, beam_width)
             st.header("Generated Code")
             for idx, code in enumerate(st.session_state.generated_codes):
                 with st.expander(f"Generated Code {idx + 1}", expanded=True):