File size: 1,946 Bytes
37197da
 
 
39ff65e
2499bf4
39ff65e
b98c632
39ff65e
 
 
 
 
 
 
37197da
39ff65e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37197da
 
 
 
39ff65e
 
 
 
 
 
37197da
 
39ff65e
37197da
 
39ff65e
 
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
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from spaces import GPU  # Import ZeroGPU decorator

# Load model & tokenizer (runs on CPU at startup)
MODEL_NAME = "ubiodee/plutus_llm"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_8bit=True
)

# Set padding token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model.eval()

# Response function with ZeroGPU decorator
@GPU
def generate_response(prompt, max_new_tokens=200, temperature=0.7, top_p=0.9):
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to("cuda")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if response.startswith(prompt):
        response = response[len(prompt):].strip()
    return response

# Gradio UI
demo = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(label="Enter your prompt", lines=4, placeholder="Ask about Plutus..."),
        gr.Slider(label="Max New Tokens", minimum=50, maximum=500, value=200, step=10),
        gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7, step=0.1),
        gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.9, step=0.05)
    ],
    outputs=gr.Textbox(label="Model Response"),
    title="Cardano Plutus AI Assistant",
    description="Ask questions about Plutus smart contracts or Cardano blockchain using ubiodee/plutus_llm."
)

if __name__ == "__main__":
    demo.launch()