|
import gradio as gr |
|
from langchain import HuggingFacePipeline, PromptTemplate, LLMChain |
|
from transformers import AutoTokenizer |
|
import transformers |
|
import torch |
|
|
|
|
|
model = "models/manohar02/NN-Llama-2-7b-finetune" |
|
|
|
|
|
pipeline = transformers.pipeline( |
|
"text-generation", |
|
model=model, |
|
torch_dtype=torch.bfloat16, |
|
max_length=20000, |
|
do_sample=True, |
|
top_k=10, |
|
num_return_sequences=1, |
|
eos_token_id=AutoTokenizer.from_pretrained(model).eos_token_id |
|
) |
|
|
|
llm = HuggingFacePipeline(pipeline=pipeline, model_kwargs={'temperature': 0}) |
|
|
|
|
|
template = """ |
|
Write a concise summary of the following text delimited by triple backquotes. |
|
'''{text}''' |
|
SUMMARY: |
|
""" |
|
|
|
prompt = PromptTemplate(template=template, input_variables=["text"]) |
|
|
|
llm_chain = LLMChain(prompt=prompt, llm=llm) |
|
|
|
|
|
def generate_summary(text): |
|
summary = llm_chain.run(text) |
|
return summary.split('SUMMARY:')[-1].strip() |
|
|
|
|
|
iface = gr.Interface(fn=generate_summary, inputs="text", outputs="text", title="LLaMA2 Summarizer") |
|
iface.launch() |