app.py

import gradio as gr
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM
from datetime import datetime
from transformers import pipeline
import re

tokenizer = AutoTokenizer.from_pretrained("ashishkgpian/best_mistral_model")
model = AutoModelForCausalLM.from_pretrained("ashishkgpian/best_mistral_model")
pipe.to('cuda')


@spaces.GPU
def model_seq_gen(prompt, model= model, tokenizer= tokenizer ,temp = 0.6) : 
        pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer)
        start = datetime.now()
        sequences = pipe(
            f'{prompt}' ,
            do_sample=True,
            max_new_tokens=100, 
            temperature=temp, 
            top_k=40, 
            top_p=0.95,
            num_return_sequences=1)
        extracted_title = sequences[0]['generated_text']
        extracted_title = re.sub(r'[\'"]', '', sequences[0]['generated_text'].split("Response :")[1])
        stop = datetime.now()
        time_taken = stop-start
        print(f"Execution Time : {time_taken}")
        return extracted_title
    



gr.Interface(
    fn=model_seq_gen,
    inputs=gr.Text(),
    outputs=gr.Gallery(),
).launch()