Update app.py
Browse files
app.py
CHANGED
@@ -141,24 +141,33 @@ def init_db():
|
|
141 |
conn.commit()
|
142 |
|
143 |
def initialize_llm():
|
144 |
-
"""Initialize the
|
145 |
try:
|
146 |
# Get API token
|
147 |
api_token = os.environ.get('HF_TOKEN')
|
148 |
if not api_token:
|
149 |
-
raise ValueError("No API token found")
|
150 |
|
151 |
-
#
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
156 |
|
157 |
-
|
158 |
-
test = llm("Say hello")
|
159 |
-
print(f"Test response: {test}")
|
160 |
|
161 |
-
return
|
162 |
|
163 |
except Exception as e:
|
164 |
print(f"LLM initialization error: {str(e)}")
|
|
|
141 |
conn.commit()
|
142 |
|
143 |
def initialize_llm():
|
144 |
+
"""Initialize the LLM using the transformers library."""
|
145 |
try:
|
146 |
# Get API token
|
147 |
api_token = os.environ.get('HF_TOKEN')
|
148 |
if not api_token:
|
149 |
+
raise ValueError("No Hugging Face API token found in environment variables.")
|
150 |
|
151 |
+
# Define model repository ID
|
152 |
+
model_name = "mistralai/Mistral-7B-Instruct-v0.3"
|
153 |
+
|
154 |
+
# Load tokenizer and model
|
155 |
+
print("Loading tokenizer...")
|
156 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=api_token)
|
157 |
+
|
158 |
+
print("Loading model...")
|
159 |
+
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=api_token)
|
160 |
+
|
161 |
+
# Test the model with a simple prompt
|
162 |
+
print("Testing the model...")
|
163 |
+
prompt = "Say hello"
|
164 |
+
inputs = tokenizer(prompt, return_tensors="pt")
|
165 |
+
outputs = model.generate(**inputs, max_length=50)
|
166 |
+
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
167 |
|
168 |
+
print(f"Test response: {response}")
|
|
|
|
|
169 |
|
170 |
+
return model, tokenizer
|
171 |
|
172 |
except Exception as e:
|
173 |
print(f"LLM initialization error: {str(e)}")
|