Update app.py

app.py

@@ -1,16 +1,24 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+import torch
 
-# Load a suitable LLM model (replace with the chosen model name)
-model_name = "databricks/dolly-v2-3b"  # Example: you can replace with a financial model
+
+# Load the Dolly model and tokenizer
+model_name = "databricks/dolly-v2-3b"  
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+
+# Function to generate model response
 def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=100)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_length=200, num_return_sequences=1)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except Exception as e:
+        return f"An error occurred: {e}"
+
 
 # Set up Gradio interface
 gr.Interface(fn=generate_response, inputs="text", outputs="text").launch()