|
import gradio as gr |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
import torch |
|
|
|
|
|
model_name = "distilgpt2" |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = AutoModelForCausalLM.from_pretrained(model_name) |
|
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
model.to(device) |
|
|
|
def generate_response(prompt): |
|
|
|
inputs = tokenizer.encode(prompt, return_tensors="pt").to(device) |
|
|
|
|
|
outputs = model.generate(inputs, max_length=150, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id) |
|
|
|
|
|
response = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
return response |
|
|
|
|
|
iface = gr.Interface( |
|
fn=generate_response, |
|
inputs="text", |
|
outputs="text", |
|
title="Crypto Analysis Model", |
|
description="Enter your prompt related to Bitcoin or cryptocurrency." |
|
) |
|
|
|
|
|
iface.launch() |
|
|