AmiyendraOP commited on
Commit
2557d7b
Β·
verified Β·
1 Parent(s): a11b121

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -0
app.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ from peft import PeftModel
4
+ import gradio as gr
5
+
6
+ base_model_id = "unsloth/llama-3-8b-bnb-4bit"
7
+ adapter_model_id = "Amiyendra/llama3-legal-assistant"
8
+
9
+ tokenizer = AutoTokenizer.from_pretrained(base_model_id)
10
+ base_model = AutoModelForCausalLM.from_pretrained(
11
+ base_model_id,
12
+ device_map="auto",
13
+ torch_dtype=torch.float16,
14
+ )
15
+
16
+
17
+ model = PeftModel.from_pretrained(base_model, adapter_model_id)
18
+ model.eval()
19
+
20
+
21
+ def generate(prompt):
22
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
23
+ with torch.no_grad():
24
+ outputs = model.generate(**inputs, max_new_tokens=512)
25
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
26
+
27
+
28
+ gr.Interface(fn=generate, inputs="text", outputs="text", title="LLaMA 3 Legal Assistant").launch()