Create app.py

app.py

@@ -0,0 +1,148 @@
+import streamlit as st
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+# Load the Starcoder2 model and tokenizer
+model_name = "starcoder2"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+def code_complete(prompt, max_length=256):
+    """
+    Generate code completion suggestions for the given prompt.
+
+    Args:
+        prompt (str): The incomplete code snippet.
+        max_length (int, optional): The maximum length of the generated code. Defaults to 256.
+
+    Returns:
+        list: A list of code completion suggestions.
+    """
+    # Tokenize the input prompt
+    inputs = tokenizer.encode_plus(prompt, 
+                                    add_special_tokens=True, 
+                                    max_length=max_length, 
+                                    padding="max_length", 
+                                    truncation=True, 
+                                    return_attention_mask=True, 
+                                    return_tensors="pt")
+
+    # Generate code completion suggestions
+    outputs = model.generate(inputs["input_ids"], 
+                             attention_mask=inputs["attention_mask"], 
+                             max_length=max_length)
+
+    # Decode the generated code
+    suggestions = []
+    for output in outputs:
+        decoded_code = tokenizer.decode(output, skip_special_tokens=True)
+        suggestions.append(decoded_code)
+
+    return suggestions
+
+def code_fix(code):
+    """
+    Fix errors in the given code snippet.
+
+    Args:
+        code (str): The code snippet with errors.
+
+    Returns:
+        str: The corrected code snippet.
+    """
+    # Tokenize the input code
+    inputs = tokenizer.encode_plus(code, 
+                                    add_special_tokens=True, 
+                                    max_length=512, 
+                                    padding="max_length", 
+                                    truncation=True, 
+                                    return_attention_mask=True, 
+                                    return_tensors="pt")
+
+    # Generate corrected code
+    outputs = model.generate(inputs["input_ids"], 
+                             attention_mask=inputs["attention_mask"], 
+                             max_length=512)
+
+    # Decode the generated code
+    corrected_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return corrected_code
+
+def text_to_code(text, max_length=256):
+    """
+    Generate code from a natural language description.
+
+    Args:
+        text (str): The natural language description of the code.
+        max_length (int, optional): The maximum length of the generated code. Defaults to 256.
+
+    Returns:
+        str: The generated code.
+    """
+    # Tokenize the input text
+    inputs = tokenizer.encode_plus(text, 
+                                    add_special_tokens=True, 
+                                    max_length=max_length, 
+                                    padding="max_length", 
+                                    truncation=True, 
+                                    return_attention_mask=True, 
+                                    return_tensors="pt")
+
+    # Generate code from the input text
+    outputs = model.generate(inputs["input_ids"], 
+                             attention_mask=inputs["attention_mask"], 
+                             max_length=max_length)
+
+    # Decode the generated code
+    generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return generated_code
+
+# Create a Streamlit app
+st.title("Codebot")
+st.write("Welcome to the Codebot! You can use this app to generate code completions, fix errors in your code, or generate code from a natural language description.")
+
+# Create a tab for code completion
+code_completion_tab = st.tab("Code Completion")
+
+with code_completion_tab:
+    st.write("Enter an incomplete code snippet:")
+    prompt_input = st.text_input("Prompt:", value="")
+    generate_button = st.button("Generate Completions")
+
+    if generate_button:
+        completions = code_complete(prompt_input)
+        st.write("Code completions:")
+        for i, completion in enumerate(completions):
+            st.write(f"{i+1}. {completion}")
+
+# Create a tab for code fixing
+code_fixing_tab = st.tab("Code Fixing")
+
+with code_fixing_tab:
+    st.write("Enter a code snippet with errors:")
+    code_input = st.text_area("Code:", height=300)
+    fix_button = st.button("Fix Errors")
+
+    if fix_button:
+        corrected_code = code_fix(code_input)
+        st.write("Corrected code:")
+        st.code(corrected_code)
+
+# Create a tab for text-to-code
+text_to_code_tab = st.tab("Text-to-Code")
+
+with text_to_code_tab:
+    st.write("Enter a natural language description of the code:")
+    text_input = st.text_input("Description:", value="")
+    generate_button = st.button("Generate Code")
+
+    if generate_button:
+        generated_code = text_to_code(text_input)
+        st.write("Generated code:")
+        st.code(generated_code)
+
+# Run the Streamlit app
+if __name__ == "__main__":
+    st.run()