Create app.py

app.py

@@ -0,0 +1,153 @@
+import streamlit as st
+import json
+from peft import AutoPeftModelForCausalLM
+from transformers import GenerationConfig, AutoTokenizer
+import torch
+import re
+
+
+def process_data_sample(example):
+    processed_example = "You have to generate api developer documentation json object which helps the user to create api documentation. ### Instruction : " + example["Instruction"] + ". ### Response :"
+    return processed_example
+
+def processing_ouput(model_response):
+    pattern = r"### Response :(.*)"
+    
+    # Find the first match
+    match = re.search(pattern, model_response, re.DOTALL)
+    
+    # Extract and process the match
+    extracted_response = ''
+    if match:
+        extracted_response = match.group(1).strip()
+    
+    print(extracted_response)
+
+    original_json_str = extracted_response
+    original_json = json.loads(original_json_str)
+
+    # New JSON structure
+    new_json = {
+    "Name": "API Name will come here",
+    "Endpoint": original_json["Endpoint"],
+    "Method": original_json["Method"],
+    "Description": original_json["Description"],
+    "Headers": original_json["Headers"],
+    "Request_Body": {
+        # Assuming this information needs to be manually added or transformed
+        "ProductID": "Unique identifier of the product for which the price is to be updated.",
+        "NewPrice": "New price to be set for the specified product."
+    },
+    "Response_Body": original_json["Response_Body"],
+    "Steps_To_Use": original_json["Steps_To_Use"],
+    "Edge_Cases": {
+        # Assuming this information needs to be manually added or transformed
+        "Invalid_ProductID": "If the specified product ID is invalid, an error message will be returned.",
+        "Negative_Price": "If the new price is negative, an error message will be returned."
+    },
+    "Exceptions": original_json["Exceptions"],
+    "Usage_Examples": original_json["Usage_Examples"]
+    }
+    # Convert the new JSON object to a string
+    new_json_str = json.dumps(new_json, indent=4)
+    return new_json_str
+
+    
+
+def model_function(input_data):
+    model_response = input_data
+    tokenizer = AutoTokenizer.from_pretrained("Shubhang999/shu3")
+
+    inp_str = process_data_sample(
+        {
+            "Instruction": input_data,
+        }
+    )
+    inputs = tokenizer(inp_str, return_tensors="pt").to("cuda")
+    
+    model = AutoPeftModelForCausalLM.from_pretrained(
+        "Shubhang999/shu3",
+        low_cpu_mem_usage=True,
+        return_dict=True,
+        torch_dtype=torch.float16,
+        device_map="cuda")
+    
+    generation_config = GenerationConfig(
+        do_sample=True,
+        top_k=1,
+        temperature=0.1,
+        max_new_tokens=800,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    outputs = model.generate(**inputs, generation_config=generation_config)
+    #print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+    processed_ouput = processing_ouput(outputs)
+    return processed_ouput
+    
+
+# Streamlit UI code
+def key_value_input(container, label, key, value, index):
+    common_headers = [
+        "Host", "User-Agent", "Accept", "Accept-Language", 
+        "Accept-Encoding", "Connection", "Referer", "Cookie", 
+        "Authorization", "Cache-Control", "Content-Type"
+    ]
+    col1, col2 = container.columns(2)
+
+    # Adding a selectbox for common headers
+    key_input = col1.selectbox(f"{label} Key {index}", options=[''] + common_headers, index=common_headers.index(key) if key in common_headers else 0, key=f"{label}_key_{index}")
+    value_input = col2.text_input(f"{label} Value {index}", value, key=f"{label}_value_{index}")
+
+    return key_input, value_input
+
+def dynamic_key_value_pairs(label):
+    container = st.container()
+    all_pairs = []
+
+    if label not in st.session_state:
+        st.session_state[label] = [{'key': '', 'value': ''}]
+
+    for i, pair in enumerate(st.session_state[label]):
+        key, value = key_value_input(container, label, pair['key'], pair['value'], i)
+        all_pairs.append((key, value))
+
+    if container.button(f"Add more to {label}"):
+        st.session_state[label].append({'key': '', 'value': ''})
+
+    return all_pairs
+
+# Streamlit UI layout
+st.title('API Documentation Generator')
+
+# Text input for API Endpoint
+api_endpoint = st.text_input("API Endpoint", "https://example.com/api")
+
+# Dropdown for API Method
+api_methods = ["GET", "POST", "PUT", "DELETE", "PATCH"]
+api_method = st.selectbox("API Method", api_methods)
+
+# Dynamic key-value pairs for Request Header, Body, and Response Object
+request_header_pairs = dynamic_key_value_pairs("Request Header")
+request_body_pairs = dynamic_key_value_pairs("Request Body")
+response_object_pairs = dynamic_key_value_pairs("Response Object")
+
+# Button to Generate Documentation
+if st.button('Generate Documentation'):
+    request_header = {k: v for k, v in request_header_pairs if k and v}
+    request_body = {k: v for k, v in request_body_pairs if k and v}
+    response_object = {k: v for k, v in response_object_pairs if k and v}
+
+    user_input = {
+        "API_Endpoint": api_endpoint,
+        "API_Method": api_method,
+        "Request_Object": request_header,
+        "Response_Object": response_object
+    }
+
+    # Call the model function with the processed input
+    documentation = model_function(user_input)
+
+    # Display the model output on the UI
+    st.write(documentation)
+