|
|
import os |
|
|
import torch |
|
|
import gradio as gr |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
|
|
from peft import PeftModel |
|
|
|
|
|
|
|
|
HF_TOKEN = os.environ.get("HF_TOKEN") |
|
|
|
|
|
|
|
|
base_model = "mistralai/Mistral-7B-v0.1" |
|
|
adapter_path = "./mistral-recovery-model" |
|
|
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(base_model, use_auth_token=HF_TOKEN) |
|
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
base_model, |
|
|
torch_dtype=torch.float32, |
|
|
use_auth_token=HF_TOKEN |
|
|
) |
|
|
|
|
|
model = PeftModel.from_pretrained(model, adapter_path, use_auth_token=HF_TOKEN) |
|
|
|
|
|
|
|
|
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.7) |
|
|
|
|
|
|
|
|
def generate(prompt): |
|
|
return pipe(prompt)[0]['generated_text'] |
|
|
|
|
|
gr.Interface( |
|
|
fn=generate, |
|
|
inputs="text", |
|
|
outputs="text", |
|
|
title="ποΈ Mistral Recovery Coach (CPU Mode)", |
|
|
description="Enter your workout summary and get personalized recovery advice." |
|
|
).launch() |
|
|
|