heyIamUmair commited on
Commit
d3f1de3
·
verified ·
1 Parent(s): 0dd3316

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -0
app.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
3
+
4
+ model_id = "heyIamUmair/llama3-3b-merged-legal"
5
+
6
+ tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
7
+ model = AutoModelForCausalLM.from_pretrained(
8
+ model_id,
9
+ device_map="auto",
10
+ torch_dtype="auto"
11
+ )
12
+
13
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
14
+
15
+ def chat(query):
16
+ output = pipe(query, max_new_tokens=200, do_sample=True, temperature=0.7)
17
+ return output[0]["generated_text"]
18
+
19
+ gr.ChatInterface(fn=chat, title="🧑‍⚖️ Pakistan Law Chatbot (LLama 3.2)", theme="default").launch()