Spaces:
Sleeping
Sleeping
File size: 2,919 Bytes
77e5700 d3ee0d7 77e5700 de3ddde 77e5700 1cfea7f 77e5700 de3ddde 77e5700 b0c41b9 77e5700 de3ddde 77e5700 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel
base_model_id = "NousResearch/Llama-2-7b-chat-hf"
lora_path = "Arsh014/lora-llama2-finetuned"
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
model = AutoModelForCausalLM.from_pretrained(
base_model_id,
load_in_8bit=True,
torch_dtype=torch.float16,
device_map="auto"
)
try:
model = PeftModel.from_pretrained(model, lora_path)
model.eval() # Set model to evaluation mode
except Exception as e:
print(f"Error loading LoRA adapter from {lora_path}. Ensure it exists and is correct.")
print(f"Error: {e}")
# 5. Create a text-generation pipeline
print("Creating text-generation pipeline.")
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer
)
def format_prompt(instruction, code):
"""Formats the instruction and input code into the required prompt template."""
return f"""### Instruction:
{instruction}
### Input:
{code}
### Response:"""
def explain_dockerfile(instruction, code):
"""Generates the explanation using the text-generation pipeline."""
if not instruction or not code:
return "Please provide both an instruction and the Dockerfile code."
prompt = format_prompt(instruction, code)
# Generate response
response = pipe(
prompt,
max_new_tokens=256,
do_sample=True,
temperature=0.7,
return_full_text=False
)
generated_text = response[0]["generated_text"].strip()
if "### Response:" in generated_text:
return generated_text.split("### Response:")[-1].strip()
return generated_text
# 6. Gradio Interface
print("Launching Gradio Interface...")
iface = gr.Interface(
fn=explain_dockerfile,
inputs=[
gr.Textbox(
label="Instruction",
placeholder="e.g., Explain the function of each line and the overall goal of this Dockerfile.",
value="Explain this Dockerfile in detail and suggest one security improvement.",
lines=2
),
gr.Textbox(
label="Dockerfile Code",
lines=10,
placeholder="Paste your Dockerfile here, e.g., \nFROM python:3.9-slim\nWORKDIR /app\nCOPY requirements.txt .\nRUN pip install -r requirements.txt\nCOPY . .\nCMD [\"python\", \"app.py\"]",
value="FROM node:18-alpine\nWORKDIR /usr/src/app\nCOPY package*.json ./ \nRUN npm install\nCOPY . .\nEXPOSE 3000\nCMD [ \"npm\", \"start\" ]"
)
],
outputs=gr.Textbox(
label="Explanation (Generated by LoRA Model)",
lines=15
),
title="LoRA-Tuned Llama-2 Dockerfile Explainer",
description="A simple application to explain complex Dockerfiles using a fine-tuned Llama-2 model (via LoRA).",
live=False
)
if __name__ == "__main__":
iface.launch()
|