back to streamlit

app.py

@@ -1,4 +1,4 @@
-import gradio as gr
+import streamlit as st
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
@@ -15,6 +15,7 @@ openai_api_key = os.environ.get('OPENAI_API_KEY')
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+@st.cache_resource
 def load_model():
     base_model = AutoModelForCausalLM.from_pretrained(base_model_id, use_auth_token=hf_token)
     model = PeftModel.from_pretrained(base_model, model_id, use_auth_token=hf_token).to(device)
@@ -72,40 +73,31 @@ def process_prompt(tokenizer, content, video_summary = '', guidelines = None):
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     return prompt
 
-def generate_headlines(content, video_summary):
-    if not content.strip():
-        return "Please enter valid article content."
-        
-    if not video_summary.strip():
-        video_summary = ''
-        
-    prompt = process_prompt(tokenizer, content, video_summary, guidelines)
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
-    
-    headlines = []
-    for i in range(5):
-        outputs = model.generate(**inputs,
-                               max_new_tokens=60,
-                               num_return_sequences=1,
-                               do_sample=True,
-                               temperature=0.7)
-        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-        response = response.replace('"', '')
-        headlines.append(f"Headline {i+1}: {response}")
-    
-    return "\n\n".join(headlines)
 
-# Create Gradio interface
-demo = gr.Interface(
-    fn=generate_headlines,
-    inputs=[
-        gr.Textbox(label="Article Content", placeholder="Type the main content of the article here..."),
-        gr.Textbox(label="Video Summary (Optional)", placeholder="Type the summary of the video related to the article...")
-    ],
-    outputs=gr.Textbox(label="Generated Headlines"),
-    title="Article Headline Writer",
-    description="Write catchy headlines from content and video summary."
-)
+st.title("Article Headline Writer")
+st.write("Write a catchy headline from content and video summary.")
 
-if __name__ == "__main__":
-    demo.launch()
+# Inputs for content and video summary
+content = st.text_area("Enter the article content:", placeholder="Type the main content of the article here...")
+video_summary = st.text_area("Enter the summary of the article's accompanying video (optional):", placeholder="Type the summary of the video related to the article...")
+
+if st.button("Generate Headline"):
+    if content.strip():
+        if not video_summary.strip():
+            video_summary = ''
+        prompt = process_prompt(tokenizer, content, video_summary, guidelines)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
+        
+        st.write("### Generated 5 Potential Headlines:")
+        for i in range(5):
+            st.write(f"### Headline {i+1}")
+            outputs = model.generate(**inputs,
+                                   max_new_tokens=60,
+                                   num_return_sequences=1, 
+                                   do_sample=True,
+                                   temperature=0.7)
+            response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+            response = response.replace('"', '')
+            st.write(f"{response}")
+    else:
+        st.write("Please enter a valid prompt.")

app_streamlit.py → app_gradio.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
@@ -15,7 +15,6 @@ openai_api_key = os.environ.get('OPENAI_API_KEY')
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-@st.cache_resource
 def load_model():
     base_model = AutoModelForCausalLM.from_pretrained(base_model_id, use_auth_token=hf_token)
     model = PeftModel.from_pretrained(base_model, model_id, use_auth_token=hf_token).to(device)
@@ -73,31 +72,40 @@ def process_prompt(tokenizer, content, video_summary = '', guidelines = None):
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     return prompt
 
+def generate_headlines(content, video_summary):
+    if not content.strip():
+        return "Please enter valid article content."
+        
+    if not video_summary.strip():
+        video_summary = ''
+        
+    prompt = process_prompt(tokenizer, content, video_summary, guidelines)
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
+    
+    headlines = []
+    for i in range(5):
+        outputs = model.generate(**inputs,
+                               max_new_tokens=60,
+                               num_return_sequences=1,
+                               do_sample=True,
+                               temperature=0.7)
+        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        response = response.replace('"', '')
+        headlines.append(f"Headline {i+1}: {response}")
+    
+    return "\n\n".join(headlines)
 
-st.title("Article Headline Writer")
-st.write("Write a catchy headline from content and video summary.")
-
-# Inputs for content and video summary
-content = st.text_area("Enter the article content:", placeholder="Type the main content of the article here...")
-video_summary = st.text_area("Enter the summary of the article's accompanying video (optional):", placeholder="Type the summary of the video related to the article...")
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_headlines,
+    inputs=[
+        gr.Textbox(label="Article Content", placeholder="Type the main content of the article here..."),
+        gr.Textbox(label="Video Summary (Optional)", placeholder="Type the summary of the video related to the article...")
+    ],
+    outputs=gr.Textbox(label="Generated Headlines"),
+    title="Article Headline Writer",
+    description="Write catchy headlines from content and video summary."
+)
 
-if st.button("Generate Headline"):
-    if content.strip():
-        if not video_summary.strip():
-            video_summary = ''
-        prompt = process_prompt(tokenizer, content, video_summary, guidelines)
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
-        
-        st.write("### Generated 5 Potential Headlines:")
-        for i in range(5):
-            st.write(f"### Headline {i+1}")
-            outputs = model.generate(**inputs,
-                                   max_new_tokens=60,
-                                   num_return_sequences=1, 
-                                   do_sample=True,
-                                   temperature=0.7)
-            response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-            response = response.replace('"', '')
-            st.write(f"{response}")
-    else:
-        st.write("Please enter a valid prompt.")
+if __name__ == "__main__":
+    demo.launch()