lost

app.py

@@ -2,10 +2,15 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
 from PyPDF2 import PdfReader
 
+
 # Models and Tokenizers Setup
 models = {
     "Text Generator (Bloom)": {
-        "model": AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m"),
+        "model": AutoModelForCausalLM.from_pretrained(
+            "bigscience/bloom-560m",
+            device_map="auto",
+            torch_dtype="auto"
+        ),
         "tokenizer": AutoTokenizer.from_pretrained("bigscience/bloom-560m"),
     },
     "PDF Summarizer (T5)": {
@@ -32,9 +37,19 @@ def chat_with_model(model_choice, user_message, chat_history, file=None):
     model = model_info["model"]
 
     # Tokenize Input
-    inputs = tokenizer(user_message, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    # Generate Output
-    outputs = model.generate(**inputs, max_length=150, num_beams=5, early_stopping=True)
+    inputs = tokenizer(user_message, return_tensors="pt", truncation=True, max_length=256)
+    
+    # Adjust max_length and parameters for the PDF summarizer
+    max_length = 100
+    num_beams = 3
+    outputs = model.generate(
+        **inputs, 
+        max_length=max_length, 
+        num_beams=num_beams, 
+        early_stopping=True, 
+        no_repeat_ngram_size=2
+    )
+    
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     # Update Chat History
@@ -50,7 +65,7 @@ def extract_text_from_pdf(file):
 
 # Interface Setup
 def create_chat_interface():
-    with gr.Blocks(css="""
+    with gr.Blocks(css=""" 
         .chatbox {
             background-color: #f7f7f8;
             border-radius: 12px;
@@ -102,4 +117,4 @@ def create_chat_interface():
 
 if __name__ == "__main__":
     interface = create_chat_interface()
-    interface.launch()
+    interface.launch(server_name="0.0.0.0", server_port=7860)